Tuesday, July 20, 2010

Nested list comprehensions in Python

Python's nested list comprehensions often hover on the edge between elegance and obfuscation.

I made some examples to show the order of evaluation, and the results of adding/removing the nested bracket symbols.


matrix = [
[1, 2, 3],
[4, 5, 6],
[7, 8, 9],
]

# This flattens and transposes the matrix.
c = [row[i] for i in range(len(matrix)) for row in matrix]
d = []
for i in range(len(matrix)):
for row in matrix:
d.append(row[i])
assert(c == d)
print c # [1, 4, 7, 2, 5, 8, 3, 6, 9]

# This flattens the matrix.
e = [row[i] for row in matrix for i in range(len(matrix))]
f = []
for row in matrix:
for i in range(len(matrix)):
f.append(row[i])
assert(e == f)
print e # [1, 2, 3, 4, 5, 6, 7, 8, 9]

# This transposes the matrix.
a = [[row[i] for row in matrix] for i in range(len(matrix))]
b1 = []
for i in range(len(matrix)):
b2 = []
for row in matrix:
b2.append(row[i])
b1.append(b2)
assert(a == b1)
print a # [[1, 4, 7], [2, 5, 8], [3, 6, 9]]

# This returns the matrix unchanged.
g = [[row[i] for i in range(len(matrix))] for row in matrix]
h1 = []
for row in matrix:
h2 = []
for i in range(len(matrix)):
h2.append(row[i])
h1.append(h2)
assert(g == h1)
assert(g == matrix)
print g # [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

# This oddity builds up a nested list from simple lists.
k = [zip(j,j) for j in [range(i) for i in range(4)]]
m1 = []
for i in range(4):
m1.append(range(i))
m2 = []
for j in m1:
m2.append(zip(j,j))
assert(k == m2)
print k # [[], [(0, 0)], [(0, 0), (1, 1)], [(0, 0), (1, 1), (2, 2)]]


Thanks to Luka Marinko for his instructions for setting up syntax highlighting on Blogger.