Skip to content

Instantly share code, notes, and snippets.

@siscia
Last active October 24, 2024 09:53
Show Gist options
  • Save siscia/988bf4523918345a6a8285b32e685e03 to your computer and use it in GitHub Desktop.
Save siscia/988bf4523918345a6a8285b32e685e03 to your computer and use it in GitHub Desktop.
Splitting a go array or slice in a defined number of chunks.
package foo
// SplitSlice splits a slice in `numberOfChunks` slices.
//
// Each slice is, at most, one element bigger than any other slice.
// Abs( len(result[i]) - len(result[j]) ) <= 1 for every i, j in the result.
//
// Moreover the function maintains the order of the elements, it is stable.
//
// If the input array is nil, or empty, the function returns nil.
// If the number of split is less than zero, the function returns nil.
// If there are more splits than elements in the input array, the function will
// return some nil slice in the result.
func SplitSlice(array []int, numberOfChunks int) [][]int {
if len(array) == 0 {
return nil
}
if numberOfChunks <= 0 {
return nil
}
if numberOfChunks == 1 {
return [][]int{array}
}
result := make([][]int, numberOfChunks)
// we have more splits than elements in the input array.
if numberOfChunks > len(array) {
for i := 0; i < len(array); i++ {
result[i] = []int{array[i]}
}
return result
}
for i := 0; i < numberOfChunks; i++ {
min := (i * len(array) / numberOfChunks)
max := ((i + 1) * len(array)) / numberOfChunks
result[i] = array[min:max]
}
return result
}
package foo
import (
"sort"
"testing"
"testing/quick"
)
func TestSplitArray(t *testing.T) {
f := func(array []int, numberOfChunks uint8) bool {
t.Logf("Running with: Len %d and Split: %d", len(array), numberOfChunks)
result := SplitSlice(array, int(numberOfChunks))
t.Logf("Result: %v", result)
if result == nil && (len(array) == 0 || numberOfChunks <= 0) {
return true
}
if result == nil {
t.Logf("Got nil result")
return false
}
if int(numberOfChunks) != len(result) {
t.Logf("Got array of wrong size. Want size: %d, Got size: %d", numberOfChunks, len(result))
return false
}
var chunkLens []int
for _, chunk := range result {
chunkLens = append(chunkLens, len(chunk))
}
totalLen := 0
for _, l := range chunkLens {
totalLen += l
}
if totalLen != len(array) {
t.Logf("Wrong total len. Want: %d, Got %d", len(array), totalLen)
return false
}
sort.Ints(chunkLens)
if chunkLens[0] > chunkLens[len(chunkLens)-1]+1 {
t.Logf("Wrong length of the bucket, found buckets with more than 2 elements of difference")
return false
}
var joinedResult []int
for _, chunk := range result {
joinedResult = append(joinedResult, chunk...)
}
for i := range joinedResult {
if joinedResult[i] != array[i] {
t.Logf("The order of elements is not respected. On index %d found different elements %d != %d", i, joinedResult[i], array[i])
}
}
return true
}
if err := quick.Check(f, nil); err != nil {
t.Error(err)
}
}
@nelz9999
Copy link

You can reduce the amount of times the math gets done:

	prev := 0
	for i := 0; i < numberOfChunks; i++ {
		max := ((i + 1) * len(array)) / numberOfChunks
		result[i] = array[prev:max]
		prev = max
	}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment