Last active
October 24, 2024 09:53
-
-
Save siscia/988bf4523918345a6a8285b32e685e03 to your computer and use it in GitHub Desktop.
Splitting a go array or slice in a defined number of chunks.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 | |
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
You can reduce the amount of times the math gets done: