This repository has been archived by the owner on Aug 25, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 2
/
zip.go
71 lines (60 loc) · 1.59 KB
/
zip.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
package iter
import "github.com/barweiss/go-tuple"
// Zip returns an iterator that yields tuples of the two provided input
// iterators.
func Zip[T, U any](a Iter[T], b Iter[U]) Iter[tuple.T2[T, U]] {
return Iter[tuple.T2[T, U]](func() (tuple.T2[T, U], bool) {
nextA, ok := a()
if !ok {
var z tuple.T2[T, U]
return z, false
}
nextB, ok := b()
if !ok {
var z tuple.T2[T, U]
return z, false
}
return tuple.New2(nextA, nextB), true
})
}
// Enumerate returns an iterator of tuples indices and values from the input
// iterator.
func Enumerate[T any](i Iter[T]) Iter[tuple.T2[int, T]] {
return Zip(Ints[int](), i)
}
// Unzip returns two iterators, one yielding the left values of the tuples
// yielded by the input iterator, the other yielding the right values of
// the tuples. Note that, while the input iterator is evaluated lazily,
// exceptionally inequal consumption of the left vs the right iterator can lead
// to high memory consumption by values cached for the other iterator.
func Unzip[T, U any](i Iter[tuple.T2[T, U]]) (Iter[T], Iter[U]) {
var aCache []T
var bCache []U
return func() (T, bool) {
if len(aCache) == 0 {
next, ok := i()
if !ok {
var z T
return z, false
}
bCache = append(bCache, next.V2)
return next.V1, true
}
res := aCache[0]
aCache = aCache[1:]
return res, true
}, func() (U, bool) {
if len(bCache) == 0 {
next, ok := i()
if !ok {
var z U
return z, false
}
aCache = append(aCache, next.V1)
return next.V2, true
}
res := bCache[0]
bCache = bCache[1:]
return res, true
}
}