Fix NSet.SetBits in GetDifference

nset.go

@@ -165,6 +165,8 @@ func (n *NSet[T]) GetBitMask(x T) StorageType {
 	return 1 << FastModPower2(((x<<BucketIndexingBits)>>BucketIndexingBits), StorageTypeBits)
 }
 
+// Union does n1=Union(n1, n2), so the current set will be updated
+// such that its a union of its old value and the passed set
 func (n *NSet[T]) Union(otherSet *NSet[T]) {
 
 	for i := 0; i < BucketCount; i++ {
@@ -192,6 +194,8 @@ func (n *NSet[T]) Union(otherSet *NSet[T]) {
 	}
 }
 
+// GetIntersection returns a new set that's the intersection between
+// this set and the passed set
 func (n *NSet[T]) GetIntersection(otherSet *NSet[T]) *NSet[T] {
 
 	outSet := NewNSet[T]()
@@ -225,6 +229,47 @@ func (n *NSet[T]) GetIntersection(otherSet *NSet[T]) *NSet[T] {
 	return outSet
 }
 
+// GetDifference returns a new set that contains the elements in this set
+// that are not in the passed set.
+//
+// For example, if s1=(1,2,3,4,5) and s2=(1,3,4), the output is
+// s3=Diff(s1,s2)=(2,5)
+func (n *NSet[T]) GetDifference(otherSet *NSet[T]) *NSet[T] {
+
+	outSet := NewNSet[T]()
+
+	for i := 0; i < BucketCount; i++ {
+
+		b1 := &n.Buckets[i]
+		b2 := &otherSet.Buckets[i]
+
+		outSet.StorageUnitCount += b1.StorageUnitCount
+
+		newB := &outSet.Buckets[i]
+		newB.StorageUnitCount = b1.StorageUnitCount
+		newB.Data = make([]StorageType, newB.StorageUnitCount)
+
+		for j := uint32(0); j < b1.StorageUnitCount && j < b2.StorageUnitCount; j++ {
+
+			newStorage := b1.Data[j] & (^b2.Data[j])
+			newB.Data[j] = newStorage
+			outSet.SetBits += uint64(bits.OnesCount64(uint64(newStorage)))
+		}
+
+		if b1.StorageUnitCount > b2.StorageUnitCount {
+
+			copy(newB.Data[b2.StorageUnitCount:], b1.Data[b2.StorageUnitCount:])
+			for j := uint32(b2.StorageUnitCount); j < newB.StorageUnitCount; j++ {
+
+				storage := newB.Data[j]
+				outSet.SetBits += uint64(bits.OnesCount64(uint64(storage)))
+			}
+		}
+	}
+
+	return outSet
+}
+
 // GetAllElements returns all the added numbers added to NSet.
 //
 // NOTE: Be careful with this if you have a lot of elements in NSet because NSet is compressed while the returned array is not.

nset_test.go

@@ -112,6 +112,21 @@ func TestNSet(t *testing.T) {
 
 	n8Elements := n8.GetAllElements()
 	AllTrue(t, len(n8Elements) == 5, n8Elements[0] == 0, n8Elements[1] == 1, n8Elements[2] == 55, n8Elements[3] == 1000, n8Elements[4] == 10000)
+
+	// GetDifference
+	nDiff1 := nset.NewNSet[uint32]()
+	nDiff1.AddMany(1, 2, 3)
+
+	nDiff2 := nset.NewNSet[uint32]()
+	nDiff2.AddMany(1, 3, 4)
+
+	nDiff3 := nDiff1.GetDifference(nDiff2)
+
+	AllTrue(t, nDiff3.SetBits == 1, nDiff3.StorageUnitCount == 1, nDiff3.Contains(2), !nDiff3.ContainsAny(1, 3, 4))
+
+	nDiff1.AddMany(1, 2, 3, 4, 5, math.MaxUint32)
+	nDiff3 = nDiff1.GetDifference(nDiff2)
+	AllTrue(t, nDiff3.SetBits == 3, nDiff3.ContainsAll(2, 5, math.MaxUint32), !nDiff3.ContainsAny(1, 3, 4))
 }
 
 func TestNSetFullRange(t *testing.T) {
@@ -151,16 +166,14 @@ func TestNSetFullRange(t *testing.T) {
 
 func AllTrue(t *testing.T, values ...bool) (success bool) {
 
-	success = true
-
 	for i := 0; i < len(values); i++ {
 		if !values[i] {
 			t.Fatalf("Expected 'true' but got 'false'\n")
-			success = false
+			return false
 		}
 	}
 
-	return success
+	return true
 }
 
 func IsEq[T comparable](t *testing.T, expected, val T) bool {