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.

Answer 1

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)]