sets.go

package sets

import "sort"

type IDSet []uint32

func Merge(a, b IDSet) (res IDSet) {
	m := make(map[uint32]uint32, len(a)+len(b))
	for _, x := range append(a, b...) {
		m[x] = 1
	}
	res = make(IDSet, len(m))
	for x := range m {
		res = append(res, x)
	}
	sort.Slice(res, func(i, j int) bool { return res[i] < res[j] })
	return res
}

func Intersect(a, b IDSet) (res IDSet) {
	if len(a) < len(b) {
		a, b = b, a
	}
	m := make(map[uint32]uint32, len(a))
	for _, x := range a {
		m[x] = 1
	}
	for _, x := range b {
		if _, ok := m[x]; ok {
			m[x]++
		}
	}
	res = make(IDSet, len(m))
	for x, n := range m {
		if n == 2 {
			res = append(res, x)
		}
	}
	sort.Slice(res, func(i, j int) bool { return res[i] < res[j] })
	return res
}

func Minus(a, b IDSet) (res IDSet) {
	if len(b) == 0 {
		res = make(IDSet, len(a))
		copy(res, a)
		return res
	}
	res = make(IDSet, len(a))
	i, j := 0, 0
	for i < len(a) && j < len(b) {
		switch {
		case b[j] < a[i]:
			j++
		case b[j] == a[i]:
			i++
			j++
		case b[j] > a[i]:
			res = append(res, a[i])
			i++
		}
	}
	for ; i < len(a); i++ {
		res = append(res, a[i])
	}
	return res
}

func IntersectMany(sets ...IDSet) (res IDSet) {
	if len(sets) == 0 {
		return nil
	}

	res = make(IDSet, len(sets[0]))
	copy(res, sets[0])
	if len(sets) > 1 {
		sort.Slice(sets, func(i, j int) bool { return len(sets[i]) < len(sets[j]) })
	}

	for i := 1; i < len(sets); i++ {
		res = Intersect(res, sets[i])
	}
	return res
}