Python list comprehension with complex data structures
Question:
I’m trying to flatten some mixed arrays in Python using LC. I’m having some trouble figuring out how to structure it.
Here’s the array’s i’m trying to flatten
arr_1 = [1, [2, 3], 4, 5]
arr_2 = [1,[2,3],[[4,5]]]
I tried this methods for arr_1 but get "TypeError: ‘int’ object is not iterable"
print([item if type(items) is list else items for items in arr_1 for item in items])
So I decided to break it into parts to see where it’s failing by using this
def check(item):
return item;
print([check(item) if type(items) is list else check(items) for items in [1, [2, 3], 4, 5] for items in arr_2])
Through the debugger I found that it’s failing at the 2d array in
for items in [1, [2, 3], 4, 5]
I don’t need the LC to be in one line but I just wanted to know how to do it in a single nested LC if its even possible.
Answers:
Using an internal stack and iter
‘s second form to simulate a while
loop:
def flatten(obj):
return [x
for stack in [[obj]]
for x, in iter(lambda: stack and [stack.pop()], [])
if isinstance(x, int)
or stack.extend(reversed(x))]
print(flatten([1, [2, 3], 4, 5]))
print(flatten([1, [2, 3], [[4, 5]]]))
print(flatten([1, [2, [], 3], [[4, 5]]]))
Output (Try it online!):
[1, 2, 3, 4, 5]
[1, 2, 3, 4, 5]
[1, 2, 3, 4, 5]
Slight variation, splitting the "long" line into two (Try it online!):
def flatten(obj):
return [x
for stack in [[obj]]
for _ in iter(lambda: stack, [])
for x in [stack.pop()]
if isinstance(x, int)
or stack.extend(reversed(x))]
To explain it a bit, here’s roughly the same with ordinary code:
def flatten(obj):
result = []
stack = [obj]
while stack:
x = stack.pop()
if isinstance(x, int):
result.append(x)
else:
stack.extend(reversed(x))
return result
If the order doesn’t matter, we can use a queue instead (inspired by 0x263A’s comment), although it’s less memory-efficient (Try it online!):
def flatten(obj):
return [x
for queue in [[obj]]
for x in queue
if isinstance(x, int) or queue.extend(x)]
We can fix the order if instead of putting each list’s contents at the end of the queue, we insert them right after the list (which is less time-efficient) in the "priority" queue (Try it online!):
def flatten(obj):
return [x
for pqueue in [[obj]]
for i, x in enumerate(pqueue, 1)
if isinstance(x, int) or pqueue.__setitem__(slice(i, i), x)]
I’m trying to flatten some mixed arrays in Python using LC. I’m having some trouble figuring out how to structure it.
Here’s the array’s i’m trying to flatten
arr_1 = [1, [2, 3], 4, 5]
arr_2 = [1,[2,3],[[4,5]]]
I tried this methods for arr_1 but get "TypeError: ‘int’ object is not iterable"
print([item if type(items) is list else items for items in arr_1 for item in items])
So I decided to break it into parts to see where it’s failing by using this
def check(item):
return item;
print([check(item) if type(items) is list else check(items) for items in [1, [2, 3], 4, 5] for items in arr_2])
Through the debugger I found that it’s failing at the 2d array in
for items in [1, [2, 3], 4, 5]
I don’t need the LC to be in one line but I just wanted to know how to do it in a single nested LC if its even possible.
Using an internal stack and iter
‘s second form to simulate a while
loop:
def flatten(obj):
return [x
for stack in [[obj]]
for x, in iter(lambda: stack and [stack.pop()], [])
if isinstance(x, int)
or stack.extend(reversed(x))]
print(flatten([1, [2, 3], 4, 5]))
print(flatten([1, [2, 3], [[4, 5]]]))
print(flatten([1, [2, [], 3], [[4, 5]]]))
Output (Try it online!):
[1, 2, 3, 4, 5]
[1, 2, 3, 4, 5]
[1, 2, 3, 4, 5]
Slight variation, splitting the "long" line into two (Try it online!):
def flatten(obj):
return [x
for stack in [[obj]]
for _ in iter(lambda: stack, [])
for x in [stack.pop()]
if isinstance(x, int)
or stack.extend(reversed(x))]
To explain it a bit, here’s roughly the same with ordinary code:
def flatten(obj):
result = []
stack = [obj]
while stack:
x = stack.pop()
if isinstance(x, int):
result.append(x)
else:
stack.extend(reversed(x))
return result
If the order doesn’t matter, we can use a queue instead (inspired by 0x263A’s comment), although it’s less memory-efficient (Try it online!):
def flatten(obj):
return [x
for queue in [[obj]]
for x in queue
if isinstance(x, int) or queue.extend(x)]
We can fix the order if instead of putting each list’s contents at the end of the queue, we insert them right after the list (which is less time-efficient) in the "priority" queue (Try it online!):
def flatten(obj):
return [x
for pqueue in [[obj]]
for i, x in enumerate(pqueue, 1)
if isinstance(x, int) or pqueue.__setitem__(slice(i, i), x)]