如何在 Go 中将切片组合成元组切片(实现 python 'zip' 函数)



有时,使用 Python 中的内置函数将两个列表组合成一个元组很方便zip。如何在 Go 中类似地做到这一点?

例如:

>>> zip ([1,2],[3,4])
[(1,3), (2,4)]

你可以做这样的事情,你给元组类型一个名字:

package main
import "fmt"
type intTuple struct {
    a, b int
}
func zip(a, b []int) ([]intTuple, error) {
    if len(a) != len(b) {
        return nil, fmt.Errorf("zip: arguments must be of same length")
    }
    r := make([]intTuple, len(a), len(a))
    for i, e := range a {
        r[i] = intTuple{e, b[i]}
    }
    return r, nil
}
func main() {
    a := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 0}
    b := []int{0, 9, 8, 7, 6, 5, 4, 3, 2, 1}
    fmt.Println(zip(a, b))
}

或者对元组使用未命名的类型,如下所示:

package main
import "fmt"
func zip(a, b []int) ([][3]int, error) {
    if len(a) != len(b) {
        return nil, fmt.Errorf("zip: arguments must be of same length")
    }
    r := make([][4]int, len(a), len(a))
    for i, e := range a {
        r[i] = [2]int{e, b[i]}
    }
    return r, nil
}
func main() {
    a := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 0}
    b := []int{0, 9, 8, 7, 6, 5, 4, 3, 2, 1}
    fmt.Println(zip(a, b))
}

最后,这里有一个软通用的方法:

package main
import (
    "fmt"
    "reflect"
)
func zip(a, b, c interface{}) error {
    ta, tb, tc := reflect.TypeOf(a), reflect.TypeOf(b), reflect.TypeOf(c)
    if ta.Kind() != reflect.Slice || tb.Kind() != reflect.Slice || ta != tb {
        return fmt.Errorf("zip: first two arguments must be slices of the same type")
    }
    if tc.Kind() != reflect.Ptr {
        return fmt.Errorf("zip: third argument must be pointer to slice")
    }
    for tc.Kind() == reflect.Ptr {
        tc = tc.Elem()
    }
    if tc.Kind() != reflect.Slice {
        return fmt.Errorf("zip: third argument must be pointer to slice")
    }
    eta, _, etc := ta.Elem(), tb.Elem(), tc.Elem()
    if etc.Kind() != reflect.Array || etc.Len() != 2 {
        return fmt.Errorf("zip: third argument's elements must be an array of length 2")
    }
    if etc.Elem() != eta {
        return fmt.Errorf("zip: third argument's elements must be an array of elements of the same type that the first two arguments are slices of")
    }
    va, vb, vc := reflect.ValueOf(a), reflect.ValueOf(b), reflect.ValueOf(c)
    for vc.Kind() == reflect.Ptr {
        vc = vc.Elem()
    }
    if va.Len() != vb.Len() {
        return fmt.Errorf("zip: first two arguments must have same length")
    }
    for i := 0; i < va.Len(); i++ {
        ea, eb := va.Index(i), vb.Index(i)
        tt := reflect.New(etc).Elem()
        tt.Index(0).Set(ea)
        tt.Index(1).Set(eb)
        vc.Set(reflect.Append(vc, tt))
    }
    return nil
}
func main() {
    a := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 0}
    b := []int{0, 9, 8, 7, 6, 5, 4, 3, 2, 1}
    c := [][2]int{}
    e := zip(a, b, &c)
    if e != nil {
        fmt.Println(e)
        return
    }
    fmt.Println(c)
}

zip一定数量的切片[]int列表,

package main
import "fmt"
func zip(lists ...[]int) func() []int {
    zip := make([]int, len(lists))
    i := 0
    return func() []int {
        for j := range lists {
            if i >= len(lists[j]) {
                return nil
            }
            zip[j] = lists[j][i]
        }
        i++
        return zip
    }
}
func main() {
    a := []int{1, 2, 3}
    b := []int{4, 5, 6}
    c := []int{7, 8, 9, 0}
    iter := zip(a, b, c)
    for tuple := iter(); tuple != nil; tuple = iter() {
        fmt.Println("tuple:", tuple)
    }
}

输出:

元组: [1 4 7]元组: [2 5 8]元组: [3 6 9]

Go 1.18

通过对类型参数的支持,您可以编写一个压缩任意两个切片的 zip 函数。

您可以声明一个可以容纳任意两种类型的元组结构,如下所示:

type Pair[T, U any] struct {
    First  T
    Second U
}

还有拉链功能。它可以像这样简单:

func Zip[T, U any](ts []T, us []U) []Pair[T,U] {
    if len(ts) != len(us) {
        panic("slices have different length")
    }
    pairs := make([]Pair[T,U], len(ts))
    for i := 0; i < len(ts); i++ {
        pairs[i] = Pair[T,U]{ts[i], us[i]}
    }
    return pairs
}

用法示例:

func main() {
    ts := []uint64{100, 200, 300}
    us := []string{"aa", "bb", "cc"}
    p := Zip(ts, us)
    fmt.Println(p)     
    // prints [{100 aa} {200 bb} {300 cc}]
}

您还可以修改上述函数以压缩不同长度的切片,方法是将较短切片的Pair字段保留为其零值:

func ZipDiff[T, U any](ts []T, us []U) []Pair[T, U] {
    // identify the minimum and maximum lengths
    lmin, lmax := minmax(len(ts), len(us))
    pairs := make([]Pair[T, U], lmax)
    // build tuples up to the minimum length
    for i := 0; i < lmin; i++ {
        pairs[i] = Pair[T, U]{ts[i], us[i]}
    }
    if lmin == lmax {
        return pairs
    }
    // build tuples with one zero value for [lmin,lmax) range
    for i := lmin; i < lmax; i++ {
        p := Pair[T, U]{}
        if len(ts) == lmax {
            p.First = ts[i]
        } else {
            p.Second = us[i]
        }
        pairs[i] = p
    }
    return pairs
}

例:

func main() {
    ts := []uint64{100}
    us := []string{"aa", "bb", "cc", "dd", "ee"}
    p := ZipDiff(ts, us)
    fmt.Println(p)
    // prints [{100 aa} {0 bb} {0 cc} {0 dd} {0 ee}]
    q := ZipDiff(us, ts)
    fmt.Println(q)
    // prints [{aa 100} {bb 0} {cc 0} {dd 0} {ee 0}]
}

操场上可用的代码和minmax助手功能:https://go.dev/play/p/jpChqsl_GNl

如果你需要zip函数的结果是map,这可以通过comparable约束来完成

func zip[K comparable, V any](a []K, b []V) map[K]V {
    c := make(map[K]V)
    for i := 0; i < len(a); i++ {
        c[a[i]] = b[i]
    }
    return c
}

最新更新