Write a function to transpose a matrix (swap its rows and columns).
Use zip(*matrix) and convert each tuple to a list, or use a nested comprehension indexing matrix[r][c]. The first row becomes the first column. The time is O(rows x cols).
Explain it simply
For a matrix with 2 rows and 3 columns, the transpose has 3 rows and 2 columns.
Line by line
def transpose(matrix):
return [list(col) for col in zip(*matrix)]
def transpose_index(matrix):
if not matrix:
return []
rows, cols = len(matrix), len(matrix[0])
return [[matrix[r][c] for r in range(rows)] for c in range(cols)]
m = [[1, 2, 3],
[4, 5, 6]]
print(transpose(m)) # Output: [[1, 4], [2, 5], [3, 6]]
print(transpose_index(m)) # Output: [[1, 4], [2, 5], [3, 6]]
print(transpose([[7]])) # Output: [[7]]*matrix unpacks the rows as separate arguments to zip, and zip pairs up the first items, then the second items, and so on. Those pairs are exactly the columns. The index version builds each new row by reading one column from every old row.
Complexity
Time O(r x c), since each element is moved once. Space O(r x c) for the new matrix.
What to say in the interview
"zip(*matrix) is the idiomatic answer: unpack the rows, zip pairs the columns, and I wrap each in list. I can also write it with index loops. It assumes all rows have equal length, because zip stops at the shortest one."