enumerators/sets.go at main · fgrzl/enumerators · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
package enumerators

// Union returns an enumerator with unique elements from both enumerators.
// Preserves order from first enumerator, then adds from second.
func Union[T comparable](first, second Enumerator[T]) Enumerator[T] {
	firstSlice, err := ToSlice(first)
	if err != nil {
		return Generate(func() (T, bool, error) {
			return *new(T), false, err
		})
	}
	secondSlice, err := ToSlice(second)
	if err != nil {
		return Generate(func() (T, bool, error) {
			return *new(T), false, err
		})
	}
	seen := make(map[T]bool)
	var result []T
	for _, item := range firstSlice {
		if !seen[item] {
			seen[item] = true
			result = append(result, item)
		}
	}
	for _, item := range secondSlice {
		if !seen[item] {
			seen[item] = true
			result = append(result, item)
		}
	}
	return Slice(result)
}

// Intersect returns an enumerator with elements present in both enumerators.
// Preserves order from first enumerator.
func Intersect[T comparable](first, second Enumerator[T]) Enumerator[T] {
	firstSlice, err := ToSlice(first)
	if err != nil {
		return Generate(func() (T, bool, error) {
			return *new(T), false, err
		})
	}
	secondSlice, err := ToSlice(second)
	if err != nil {
		return Generate(func() (T, bool, error) {
			return *new(T), false, err
		})
	}
	secondSet := make(map[T]bool)
	for _, item := range secondSlice {
		secondSet[item] = true
	}
	var result []T
	seen := make(map[T]bool)
	for _, item := range firstSlice {
		if secondSet[item] && !seen[item] {
			seen[item] = true
			result = append(result, item)
		}
	}
	return Slice(result)
}

// Except returns an enumerator with elements from first enumerator not in second.
// Preserves order from first enumerator.
func Except[T comparable](first, second Enumerator[T]) Enumerator[T] {
	firstSlice, err := ToSlice(first)
	if err != nil {
		return Generate(func() (T, bool, error) {
			return *new(T), false, err
		})
	}
	secondSlice, err := ToSlice(second)
	if err != nil {
		return Generate(func() (T, bool, error) {
			return *new(T), false, err
		})
	}
	secondSet := make(map[T]bool)
	for _, item := range secondSlice {
		secondSet[item] = true
	}
	var result []T
	for _, item := range firstSlice {
		if !secondSet[item] {
			result = append(result, item)
		}
	}
	return Slice(result)
}