Add Zip and ZipLongest function (#193)

Zip will return a new slice containing pairs with elements from input slices.
If input slices have different length, the output slice will be truncated to
the length of the smallest input slice.

ZipLongest will return a new slice containing pairs with elements from input slices.
If input slices have diffrent length, missing elements will be padded with default values.

These are the same as zip() and itertools.zip_longest() in Python.

v2/zip.go

@@ -0,0 +1,27 @@
+package pie
+
+// A pair struct containing two zipped values.
+type Zipped[T1, T2 any] struct {
+	A T1
+	B T2
+}
+
+// Zip will return a new slice containing pairs with elements from input slices.
+// If input slices have diffrent length, the output slice will be truncated to
+// the length of the smallest input slice.
+func Zip[T1, T2 any](ss1 []T1, ss2 []T2) []Zipped[T1, T2] {
+	var minLen int
+
+	if len(ss1) <= len(ss2) {
+		minLen = len(ss1)
+	} else {
+		minLen = len(ss2)
+	}
+
+	ss3 := []Zipped[T1, T2]{}
+	for i := 0; i < minLen; i++ {
+		ss3 = append(ss3, Zipped[T1, T2]{ss1[i], ss2[i]})
+	}
+
+	return ss3
+}

v2/zip_longest.go

@@ -0,0 +1,37 @@
+package pie
+
+// ZipLongest will return a new slice containing pairs with elements from input slices.
+// If input slices have diffrent length, missing elements will be padded with default values.
+func ZipLongest[T1, T2 any](ss1 []T1, ss2 []T2) []Zipped[T1, T2] {
+	var minLen, maxLen int
+	var small int8
+
+	if len(ss1) <= len(ss2) {
+		small = 1
+		minLen = len(ss1)
+		maxLen = len(ss2)
+	} else {
+		small = 2
+		minLen = len(ss2)
+		maxLen = len(ss1)
+	}
+
+	ss3 := []Zipped[T1, T2]{}
+	for i := 0; i < minLen; i++ {
+		ss3 = append(ss3, Zipped[T1, T2]{ss1[i], ss2[i]})
+	}
+
+	if small == 1 {
+		var t T1
+		for i := minLen; i < maxLen; i++ {
+			ss3 = append(ss3, Zipped[T1, T2]{t, ss2[i]})
+		}
+	} else {
+		var t T2
+		for i := minLen; i < maxLen; i++ {
+			ss3 = append(ss3, Zipped[T1, T2]{ss1[i], t})
+		}
+	}
+
+	return ss3
+}

v2/zip_longest_test.go

@@ -0,0 +1,18 @@
+package pie_test
+
+import (
+	"github.com/elliotchance/pie/v2"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestZipLongest(t *testing.T) {
+	for _, test := range zipTests {
+
+		t.Run("", func(t *testing.T) {
+			c := pie.ZipLongest(test.ss1, test.ss2)
+
+			assert.Equal(t, c, test.expectedLong)
+		})
+	}
+}

v2/zip_test.go

@@ -0,0 +1,62 @@
+package pie_test
+
+import (
+	"github.com/elliotchance/pie/v2"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+var zipTests = []struct {
+	ss1           []int
+	ss2           []float32
+	expectedShort []pie.Zipped[int, float32]
+	expectedLong  []pie.Zipped[int, float32]
+}{
+	{
+		[]int{},
+		[]float32{},
+		[]pie.Zipped[int, float32]{},
+		[]pie.Zipped[int, float32]{},
+	},
+	{
+		[]int{1, 2, 3, 4, 5},
+		[]float32{},
+		[]pie.Zipped[int, float32]{},
+		[]pie.Zipped[int, float32]{{1, 0}, {2, 0}, {3, 0}, {4, 0}, {5, 0}},
+	},
+	{
+		[]int{},
+		[]float32{1.0, 2.0, 3.0, 4.0, 5.0},
+		[]pie.Zipped[int, float32]{},
+		[]pie.Zipped[int, float32]{{0, 1.0}, {0, 2.0}, {0, 3.0}, {0, 4.0}, {0, 5.0}},
+	},
+	{
+		[]int{1, 2, 3, 4, 5},
+		[]float32{1.0, 2.0, 3.0, 4.0, 5.0},
+		[]pie.Zipped[int, float32]{{1, 1.0}, {2, 2.0}, {3, 3.0}, {4, 4.0}, {5, 5.0}},
+		[]pie.Zipped[int, float32]{{1, 1.0}, {2, 2.0}, {3, 3.0}, {4, 4.0}, {5, 5.0}},
+	},
+	{
+		[]int{1, 2, 3, 4, 5, 6, 7, 8},
+		[]float32{1.0, 2.0, 3.0, 4.0, 5.0},
+		[]pie.Zipped[int, float32]{{1, 1.0}, {2, 2.0}, {3, 3.0}, {4, 4.0}, {5, 5.0}},
+		[]pie.Zipped[int, float32]{{1, 1.0}, {2, 2.0}, {3, 3.0}, {4, 4.0}, {5, 5.0}, {6, 0}, {7, 0}, {8, 0}},
+	},
+	{
+		[]int{1, 2, 3},
+		[]float32{1.0, 2.0, 3.0, 4.0, 5.0},
+		[]pie.Zipped[int, float32]{{1, 1.0}, {2, 2.0}, {3, 3.0}},
+		[]pie.Zipped[int, float32]{{1, 1.0}, {2, 2.0}, {3, 3.0}, {0, 4.0}, {0, 5.0}},
+	},
+}
+
+func TestZip(t *testing.T) {
+	for _, test := range zipTests {
+
+		t.Run("", func(t *testing.T) {
+			c := pie.Zip(test.ss1, test.ss2)
+
+			assert.Equal(t, c, test.expectedShort)
+		})
+	}
+}