zip(), zip3(), unzip(), and unzip3() for Swift

unzip.swift

// Given array of 2-tuples, return two arrays
func unzip<T, U>(array: [(T, U)]) -> ([T], [U]) {
    var t = Array<T>()
    var u = Array<U>()
    for (a, b) in array {
        t.append(a)
        u.append(b)
    }
    return (t, u)
}

// Given sequence of 2-tuples, return two arrays
func unzip<T, U>(sequence: SequenceOf<(T, U)>) -> ([T], [U]) {
    var t = Array<T>()
    var u = Array<U>()
    for (a, b) in sequence {
        t.append(a)
        u.append(b)
    }
    return (t, u)
}

// Given array of 3-tuples, return three arrays
func unzip3<T, U, V>(array: [(T, U, V)]) -> ([T], [U], [V]) {
    var t = Array<T>()
    var u = Array<U>()
    var v = Array<V>()
    for (a, b, c) in array {
        t.append(a)
        u.append(b)
        v.append(c)
    }
    return (t, u, v)
}

// Given sequence of 3-tuples, return three arrays
func unzip3<T, U, V>(sequence: SequenceOf<(T, U, V)>) -> ([T], [U], [V]) {
    var t = Array<T>()
    var u = Array<U>()
    var v = Array<V>()
    for (a, b, c) in sequence {
        t.append(a)
        u.append(b)
        v.append(c)
    }
    return (t, u, v)
}

let elements = [(0, "Zero"), (1, "One"), (2, "Two")]
let (numbers, strings) = unzip(elements)
println("numbers = \(numbers)")
println("strings = \(strings)")
// numbers = [0, 1, 2]
// strings = ["Zero", "One", "Two]


let elements3 = [(0, "Zero", "A"), (1, "One", "B"), (2, "Two", "C")]
let (numbers3, strings3, letters3) = unzip3(elements3)
println("numbers3 = \(numbers3)")
println("strings3 = \(strings3)")
println("letters3 = \(letters3)")
// numbers3 = [0, 1, 2]
// strings3 = ["Zero", "One", "Two]
// letters3 = ["A", "B", "C"]

zip.swift

// zip

// Given two sequences, return a sequence of 2-tuples (pairs)
public func zip<A: SequenceType, B: SequenceType>(a: A, b: B)
    -> ZipSequence<A, B>
{
    return ZipSequence(a, b)
}

// Lazy sequence of tuples created from values from two other sequences
public struct ZipSequence<A: SequenceType, B: SequenceType>: SequenceType {
    private var a: A
    private var b: B
    
    public init (_ a: A, _ b: B) {
        self.a = a
        self.b = b
    }
    
    public func generate() -> ZipGenerator<A.Generator, B.Generator> {
        return ZipGenerator(a.generate(), b.generate())
    }
}

// Generator that creates tuples of values from two other generators
public struct ZipGenerator<A: GeneratorType, B: GeneratorType>: GeneratorType {
    private var a: A
    private var b: B
    
    public init(_ a: A, _ b: B) {
        self.a = a
        self.b = b
    }
    
    mutating public func next() -> (A.Element, B.Element)? {
        switch (a.next(), b.next()) {
        case let (.Some(aValue), .Some(bValue)):
            return (aValue, bValue)
        default:
            return nil
        }
    }
}


// zip3

// Given three sequences, return a sequence of 3-tuples
public func zip3<A: SequenceType, B: SequenceType, C: SequenceType>(a: A, b: B, c: C)
    -> ZipSequence3<A, B, C>
{
    return ZipSequence3(a, b, c)
}

// Sequence of tuples created from values from three other sequences
public struct ZipSequence3<A: SequenceType, B: SequenceType, C: SequenceType>: SequenceType {
    private var a: A
    private var b: B
    private var c: C
    
    public init (_ a: A, _ b: B, _ c: C) {
        self.a = a
        self.b = b
        self.c = c
    }
    
    public func generate() -> ZipGenerator3<A.Generator, B.Generator, C.Generator> {
        return ZipGenerator3(a.generate(), b.generate(), c.generate())
    }
}

// Generator that creates tuples of values from three other generators
public struct ZipGenerator3<A: GeneratorType, B: GeneratorType, C: GeneratorType>: GeneratorType {
    private var a: A
    private var b: B
    private var c: C
    
    public init(_ a: A, _ b: B, _ c: C) {
        self.a = a
        self.b = b
        self.c = c
    }
    
    mutating public func next() -> (A.Element, B.Element, C.Element)? {
        switch (a.next(), b.next(), c.next()) {
        case let (.Some(aValue), .Some(bValue), .Some(cValue)):
            return (aValue, bValue, cValue)
        default:
            return nil
        }
    }
}


// Examples

let numbers = [0, 1, 2, 3, 4]
let names = ["Zero", "One", "Two", "Three"]
let letters = ["A", "B", "C"]

let result = zip(numbers, names)
let printableResult = ", ".join(map(result) { "(\($0.0), \($0.1))" })
println("result: \(printableResult)")
// result: (0, Zero), (1, One), (2, Two), (3, Three)

let result3 = zip3(numbers, names, letters)
let printableResult3 = ", ".join(map(result3) { "(\($0.0), \($0.1), \($0.2))" })
println("result3: \(printableResult3)")
// result3: (0, Zero, A), (1, One, B), (2, Two, C)

let resultOne = filter(result3) { $0.1 == "One" }
let printableResultOne = ", ".join(resultOne.map { "(\($0.0), \($0.1), \($0.2))" })
println("resultOne: \(printableResultOne)")
// resultOne: (1, One, B)


// Note: Standard Swift Library provides Zip2, which is a struct
// but we can use it almost like a function.
let zip2Result = Zip2(numbers, names)
let printableZip2Result = ", ".join(map(zip2Result) { "(\($0.0), \($0.1))" })
println("zip2Result: \(printableZip2Result)")
// result: (0, Zero), (1, One), (2, Two), (3, Three)

After I posted this, @practicalswift pointed out that the Swift Standard Library includes a "Zip2" struct that implements the same thing that my zip() function does, without the function. It looks like it is essentially the same as my ZipSequence struct. The Swift Standard Library also has a Zip2Generator that looks a lot like my ZipGenerator.

I did search the Standard Library for "zip" before writing these, and found nothing, but I guess I must have done a case-sensitive search.

BTW, I used the zip/unzip and zip3/unzip3 function names familiar to Haskell and F# programmers. In Swift, it probably makes more sense to call both functions "zip/unzip", as we can rely on overload resolution based upon numbers and types of parameters.

For unzip/unzip3, I provide functions that take either an array or a SequenceOf. I couldn't figure out how to declare a function that just takes a Sequence where GeneratorType.Element is a tuple. Can someone help?

@airspeedswift posted a blog entry inspired by these definitions: http://airspeedvelocity.net/2014/08/02/tuple-wrangling/