Chunk (#184)

Chunk splits the input and returns multi slices whose length equals
chunkLength, except for the last slice which may contain fewer elements.

Co-authored-by: morrisyang <[email protected]>

v2/chunk.go

@@ -0,0 +1,39 @@
+package pie
+
+// Chunk splits the input and returns multi slices whose length equals chunkLength,
+// except for the last slice which may contain fewer elements.
+//
+// Examples:
+//
+//   Chunk([1, 2, 3], 4) => [ [1, 2, 3] ]
+//   Chunk([1, 2, 3], 3) => [ [1, 2, 3] ]
+//   Chunk([1, 2, 3], 2) => [ [1, 2], [3] ]
+//   Chunk([1, 2, 3], 1) => [ [1], [2], [3] ]
+//   Chunk([], 1)        => [ [] ]
+//   Chunk([1, 2, 3], 0) => panic: chunkLength should be greater than 0
+func Chunk[T any](ss []T, chunkLength int) [][]T {
+	if chunkLength <= 0 {
+		panic("chunkLength should be greater than 0")
+	}
+
+	result := make([][]T, 0)
+	l := len(ss)
+	if l == 0 {
+		return result
+	}
+
+	var step = l / chunkLength
+	if step == 0 {
+		result = append(result, ss)
+		return result
+	}
+	var remain = l % chunkLength
+	for i := 0; i < step; i++ {
+		result = append(result, ss[i*chunkLength:(i+1)*chunkLength])
+	}
+	if remain != 0 {
+		result = append(result, ss[step*chunkLength:l])
+	}
+
+	return result
+}

v2/chunk_test.go

@@ -0,0 +1,28 @@
+package pie
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+var chunkTests = []struct {
+	ss          []int
+	chunkLength int
+	expected    [][]int
+}{
+	{nil, 1, [][]int{}},
+	{[]int{}, 1, [][]int{}},
+	{[]int{1, 2, 3}, 4, [][]int{{1, 2, 3}}},
+	{[]int{1, 2, 3}, 3, [][]int{{1, 2, 3}}},
+	{[]int{1, 2, 3}, 2, [][]int{{1, 2}, {3}}},
+	{[]int{1, 2, 3}, 1, [][]int{{1}, {2}, {3}}},
+}
+
+func TestChunk(t *testing.T) {
+	for _, test := range chunkTests {
+		t.Run("", func(t *testing.T) {
+			assert.Equal(t, test.expected, Chunk(test.ss, test.chunkLength))
+		})
+	}
+}