CODEWITHSUNDEEP
pythonpython-3.xmultidimensional-array
CameronAtkinson
CameronAtkinson

Reputation: 107

Flatten multi dimensional array in python 3

I have a list of numbers: testList = [1, [1], [12], 2, 3]

I want it to become: flatList = [1, 1, 12, 2, 3]

Using a typical list comprehension such as below is not working.

flatList = [val for sublist in testList for val in sublist]
TypeError: 'int' object is not iterable

I suspected it is because the un-nested items are being treated as iterable sublists, so I tried this:

flatList = [val if isinstance(sublist, int) == False else val for sublist in testlist for val in sublist]

But I am unclear on the syntax, or if there is some better way to do this. Trying to remove val from the else clause means val is undefined. As is, it still gives me the same TypeError.

The code below does work for me, but I am interested to see if it can be done in list comprehension style, and people's opinions on that.

for sublist in testlist:
    if type(sublist) == int:
        flat.append(sublist)
    else:
        for val in sublist:
            flat.append(val)
print(flat)

>>>[1, 1, 12, 2, 3]

Upvotes: 9

Views: 5649

Answers (2)

Eric Duminil
Eric Duminil

Reputation: 54283

Since you're using Python 3, you can take advantage of yield from with a recursive function. It has been introduced in Python 3.3.

As a bonus, you can flatten arbitrary nested lists, tuples, sets or ranges:

test_list = [1, [1], [12, 'test', set([3, 4, 5])], 2, 3, ('hello', 'world'), [range(3)]]

def flatten(something):
    if isinstance(something, (list, tuple, set, range)):
        for sub in something:
            yield from flatten(sub)
    else:
        yield something


print(list(flatten(test_list)))
# [1, 1, 12, 'test', 3, 4, 5, 2, 3, 'hello', 'world', 0, 1, 2]
print(list(flatten('Not a list')))
# ['Not a list']
print(list(flatten(range(10))))
# [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

Here's another example with a debug line:

def flatten(something, level=0):
    print("%sCalling flatten with %r" % ('  ' * level, something))
    if isinstance(something, (list, tuple, set, range)):
        for sub in something:
            yield from flatten(sub, level+1)
    else:
        yield something

list(flatten([1, [2, 3], 4]))
#Calling flatten with [1, [2, 3], 4]
#  Calling flatten with 1
#  Calling flatten with [2, 3]
#    Calling flatten with 2
#    Calling flatten with 3
#  Calling flatten with 4

Upvotes: 8

kaidokuuppa
kaidokuuppa

Reputation: 642

If the sublists always contain only one item then

flatList = [item[0] if isinstance(item, list) else item for item in testList]

Upvotes: 2

Related Questions