Get column from 2D array – how to restrict array type in extension?

Question

I'd like to extend Array in Swift to return a single element in each array, or column, for a 2D array. So far I have:

extension Array where // what goes here?
    func getColumn( column: Int ) -> [ Int ] {
        return self.map { $0[ column ] }
    }
}

I believe that I need to somehow specify a 2D array after where, but I have been unable to figure out the correct way to do that.

What is the correct syntax for specifying a 2D array after the where?

I'm also curious if there is a good documentation for how to specify what is available for after where in an extension lives. I couldn't find that at Apple's Swift extension documentation

Thanks in advance.

Answer 1

You need to constrain the Element type of the array. The subscript method is defined in the CollectionType protocol:

public protocol CollectionType : Indexable, SequenceType {
    // ...
    public subscript (position: Self.Index) -> Self.Generator.Element { get }
    // ...
}

Therefore you can define an extension method for arrays whose elements are collections:

extension Array where Element : CollectionType {
    func getColumn(column : Element.Index) -> [ Element.Generator.Element ] {
        return self.map { $0[ column ] }
    }
}

Example:

let a = [[1, 2, 3], [4, 5, 6]]
let c = a.getColumn(1)

print(c) // [2, 5]

You could even define it as an additional subscripting method:

extension Array where Element : CollectionType {
    subscript(column column : Element.Index) -> [ Element.Generator.Element ] {
        return map { $0[ column ] }
    }
}

let a = [["a", "b", "c"], [ "d", "e", "f" ]]
let c = a[column: 2]
print(c) // ["c", "f"]

---

Update for Swift 3:

extension Array where Element : Collection {
    func getColumn(column : Element.Index) -> [ Element.Iterator.Element ] {
        return self.map { $0[ column ] }
    }
}

or as subscript:

extension Array where Element : Collection {
    subscript(column column : Element.Index) -> [ Element.Iterator.Element ] {
        return map { $0[ column ] }
    }
}