collections-go/set.go

package collections

// Set is a type alias over a map to an empty struct. This allows a nice
// API over the hash map functionality already available, where you can instead
// view the set as a single dimension collection and save a lot of clutter of
// empty struct initialization any time you add to the set.
type Set[T comparable] map[T]struct{}

// Initialize a new Set by providing a slice. The Set will have the
// same type as the slice.
//
//	Time Complexity: O(n)
//	Space Complexity: O(n)
//	Allocations: 1 set, n elements
func NewSet[T comparable](slice []T) Set[T] {
	set := make(Set[T])
	for _, el := range slice {
		set[el] = struct{}{}
	}
	return set
}

// Initialize a new set while preserving the elements that were found to be duplicate.
// The duplicate slice will be nil if there are no duplicates to save on allocations.
//
//	Time Complexity: O(n)
//	Space Complexity: O(n)
//	Allocations (best case): 1 set, n elements
//	Allocations (worst case): 1 set, n/2 elements, 2 slices, one with n-1 space, the
//	other is the first one resized based on number of duplicates.
func NewSetSaveDuplicates[T comparable](slice []T) (Set[T], []T) {
	set := make(Set[T])
	var dupes []T = nil
	dupesIdx := 0
	for _, el := range slice {
		if set.Contains(el) {
			if dupes == nil {
				// Now that we know there is at least one dupe, we can make the allocation.
				// The max possible size of the duplicate array is n-1 (aka a slice where
				// every element is the same).
				dupes = make([]T, len(slice)-1)
			}
			dupes[dupesIdx] = el
			dupesIdx++
		} else {
			set[el] = struct{}{}
		}
	}
	return set, dupes[:dupesIdx]
}

// ToSlice will take all of the elements of the set and create a new slice out of them.
//
//	Time Complexity: O(n)
//	Space Complexity: O(n)
//	Allocations: 1 slice, n elements
func (set Set[T]) ToSlice() []T {
	slice := make([]T, len(set))
	i := 0
	for el := range set {
		slice[i] = el
		i++
	}
	return slice
}

// Add an element to the set.
//
//	Time Complexity: O(1)
//	Space Complexity: O(1)
//	Allocations: 1 elements
func (s Set[T]) Add(el T) {
	s[el] = struct{}{}
}

// Remove an element from the set.
//
//	Time Complexity: O(1)
//	Space Complexity: O(1)
//	Allocations: None
func (s Set[T]) Remove(el T) {
	delete(s, el)
}

// Check if the set contains a particular value.
//
//	Time Complexity: O(1)
//	Space Complexity: O(1)
//	Allocations: None
func (s Set[T]) Contains(el T) bool {
	_, containsEl := s[el]
	return containsEl
}

// Creates a Clone of the set which contains all the same values.
//
//	Time Complexity: O(n)
//	Space Complexity: O(n)
//	Allocations: 1 set, n elements
func (s Set[T]) Clone() Set[T] {
	clone := make(Set[T])
	for el := range s {
		clone.Add(el)
	}
	return clone
}

// Checks if the set is equal to another.
//
//	Time Complexity: O(n)
//	Space Complexity: O(1)
//	Allocations: None
func (s Set[T]) Equals(s2 Set[T]) bool {
	if len(s) != len(s2) {
		return false
	}
	for el := range s {
		if !s2.Contains(el) {
			return false
		}
	}
	return true
}