miércoles, 20 de diciembre de 2017

Nested list comprehension in Python

When you code in python, you will use list comprehension eventually. It's not hard, but when it comes to NESTED list comprehension, it's a bit tricky.

As I go through this over an over, I just write it down to help others.

CODE: (you can copy and paste in your IDE or IDLE)

Note: code syntax highligth does NOT match Python syntax highlight. I use the colors to make the relation between with and without list comprehension

#####################################################
# first list
a = [1, 2, 3, 4]

# second list
b = ['a', 'b']

# expected list
exp = ['a1', 'a2', 'a3', 'a4', 'b1', 'b2', 'b3', 'b4']

# WITHOUT list comprehension
result = []
for item_a in a:
    for item_b in b:
        result.append(item_b+str(item_a))

print('without', result)

# WITH list comprehension
result = [item_b+str(item_a) for item_a in a for item_b in b]

print('with', result)
######################################################

This is the easy one.. but what happens if we have a nested list, let's see..

#####################################################
# first list
a = [ [1, 2], [3, 4] ]

# second list
b = ['a', 'b']

# expected list
exp = ['a1', 'a2', 'a3', 'a4', 'b1', 'b2', 'b3', 'b4']

# WITHOUT list comprehension
result = []
for item_b in b:
    for list_a in a:
        for item_a in list_a:
            result.append(item_b+str(item_a))

print('without', result)

# WITH list comprehension
result = [item_b+str(item_a) for item_b in b for list_a in a for item_a in list_a]

print('with', result)
######################################################

I guess you can see the pattern =)

So, the structure will be:

[ OPERATION for #1 for #2 ... for #N]