Flip Matrix (Lambda School iOS11 Code Challenge)

FlipMatrix.swift

/*
Given a matrix A, return the transpose of A.
The transpose of a matrix is the matrix flipped over it's main diagonal,
switching the row and column indices of the matrix.
 
Example 1:
Input: [[1,2,3],[4,5,6],[7,8,9]]
Output: [[1,4,7],[2,5,8],[3,6,9]]
 
Example 2:
Input: [[1,2,3],[4,5,6]]
Output: [[1,4],[2,5],[3,6]]
 */

func flipMatrix(_ input: [[Int]]) -> [[Int]] {
    var output = [[Int]]()
    for i in 0 ..< input.count {
        for j in 0 ..< input[i].count {
            if output.count <= j { output.append([]) }
            output[j].append(input[i][j])
        }
    }
    return output
}

// tests
print(flipMatrix([[1,2,3],[4,5,6],[7,8,9]]))
print(flipMatrix([[1,2,3],[4,5,6]]))