Showing posts with label list comprehension. Show all posts
Showing posts with label list comprehension. Show all posts

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.