-
Notifications
You must be signed in to change notification settings - Fork 129
/
slices.go
101 lines (91 loc) · 2.26 KB
/
slices.go
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
92
93
94
95
96
97
98
99
100
101
package slices
import (
"fmt"
"math"
goslices "golang.org/x/exp/slices"
)
// PartitionToLen partitions the elements of s into non-overlapping slices,
// such that each such slice contains at most maxLen elements.
func PartitionToMaxLen[S ~[]E, E any](s S, maxLen int) []S {
n := int(math.Ceil(float64(len(s)) / float64(maxLen)))
if n == 0 {
n = 1
}
return Partition(s, n)
}
// Partition partitions the elements of s into n non-overlapping slices,
// such that some slices have len(s)/n+1 items and some len(s)/n items.
// Ordering is preserved, such that Flatten(Partition(s)) is equal to s.
func Partition[S ~[]E, E any](s S, n int) []S {
if n < 1 {
panic(fmt.Sprintf("n is %d but must be at least 1", n))
}
k := len(s) - (len(s)/n)*n
rv := make([]S, n)
i := 0
for j := 0; j < k; j++ {
rv[j] = goslices.Clone(s[i : i+len(s)/n+1])
i += len(s)/n + 1
}
for j := k; j < n; j++ {
rv[j] = goslices.Clone(s[i : i+len(s)/n])
i += len(s) / n
}
return rv
}
// Flatten merges a slice of slices into a single slice.
func Flatten[S ~[]E, E any](s []S) S {
n := 0
allNil := true
for _, si := range s {
n += len(si)
allNil = allNil && si == nil
}
if allNil {
return nil
}
rv := make(S, n)
i := 0
for _, si := range s {
for _, e := range si {
rv[i] = e
i++
}
}
return rv
}
// GroupByFunc groups the elements e_1, ..., e_n of s into separate slices by keyFunc(e).
func GroupByFunc[S ~[]E, E any, K comparable](s S, keyFunc func(E) K) map[K]S {
rv := make(map[K]S)
for _, e := range s {
k := keyFunc(e)
rv[k] = append(rv[k], e)
}
return rv
}
// MapAndGroupByFuncs groups the elements e_1, ..., e_n of s into separate slices by keyFunc(e)
// and then maps those resulting elements by mapFunc(e).
func MapAndGroupByFuncs[S ~[]E, E any, K comparable, V any](s S, keyFunc func(E) K, mapFunc func(E) V) map[K][]V {
rv := make(map[K][]V)
for _, e := range s {
k := keyFunc(e)
rv[k] = append(rv[k], mapFunc(e))
}
return rv
}
func Subtract[T comparable](list []T, toRemove []T) []T {
if list == nil {
return nil
}
out := make([]T, 0, len(list))
toRemoveMap := make(map[T]bool, len(toRemove))
for _, val := range toRemove {
toRemoveMap[val] = true
}
for _, val := range list {
if !toRemoveMap[val] {
out = append(out, val)
}
}
return out
}