Skip to content

Commit

Permalink
fix: order of GroupBy and PartitionBy
Browse files Browse the repository at this point in the history
  • Loading branch information
liyishuai committed Jan 7, 2025
1 parent 5044a1d commit d3b5701
Showing 1 changed file with 18 additions and 39 deletions.
57 changes: 18 additions & 39 deletions parallel/slice.go
Original file line number Diff line number Diff line change
Expand Up @@ -66,65 +66,44 @@ func Times[T any](count int, iteratee func(index int) T) []T {
}

// GroupBy returns an object composed of keys generated from the results of running each element of collection through iteratee.
// The order of grouped values is determined by the order they occur in the collection.
// `iteratee` is call in parallel.
func GroupBy[T any, U comparable, Slice ~[]T](collection Slice, iteratee func(item T) U) map[U]Slice {
result := map[U]Slice{}

var mu sync.Mutex
var wg sync.WaitGroup
wg.Add(len(collection))

for _, item := range collection {
go func(_item T) {
key := iteratee(_item)

mu.Lock()

result[key] = append(result[key], _item)
keys := Map(collection, func(item T, _ int) U {
return iteratee(item)
})

mu.Unlock()
wg.Done()
}(item)
for i, item := range collection {
result[keys[i]] = append(result[keys[i]], item)
}

wg.Wait()

return result
}

// PartitionBy returns an array of elements split into groups. The order of grouped values is
// determined by the order they occur in collection. The grouping is generated from the results
// of running each element of collection through iteratee.
// The order of groups is determined by their first appearance in the collection.
// `iteratee` is call in parallel.
func PartitionBy[T any, K comparable, Slice ~[]T](collection Slice, iteratee func(item T) K) []Slice {
result := []Slice{}
seen := map[K]int{}

var mu sync.Mutex
var wg sync.WaitGroup
wg.Add(len(collection))

for _, item := range collection {
go func(_item T) {
key := iteratee(_item)

mu.Lock()

resultIndex, ok := seen[key]
if !ok {
resultIndex = len(result)
seen[key] = resultIndex
result = append(result, []T{})
}

result[resultIndex] = append(result[resultIndex], _item)
keys := Map(collection, func(item T, _ int) K {
return iteratee(item)
})

mu.Unlock()
wg.Done()
}(item)
for i, item := range collection {
if resultIndex, ok := seen[keys[i]]; ok {
result[resultIndex] = append(result[resultIndex], item)
} else {
resultIndex = len(result)
seen[keys[i]] = resultIndex
result = append(result, Slice{item})
}
}

wg.Wait()

return result
}

0 comments on commit d3b5701

Please sign in to comment.