CODEWITHSUNDEEP
pythonpython-3.xlist-comprehension
cardamom
cardamom

Reputation: 7421

list comprehension on elements of varying length

Some of the elements in my list are a bit longer and need different treatment. This list looks like this:

a = [['red', 'square', 'up'], ['red', 'circle'], 
     ['blue', 'triangle'], ['blue', 'square'], 
     ['blue', 'octagon'], ['blue', 'diamond', 'down']]

I can almost get it do do what I want like this:

[[x[0], x[1] + 's'] for x in a]

which gives

[['red', 'squares'],
 ['red', 'circles'],
 ['blue', 'triangles'],
 ['blue', 'squares'],
 ['blue', 'octagons'],
 ['blue', 'diamonds']]

but that is losing information from the longer elements. The desired output is:

[['red', 'squares', 'up1'], ['red', 'circles'], 
['blue', 'triangles'], ['blue', 'squares'], 
 ['blue', 'octagons'], ['blue', 'diamonds', 'down1']]

Just doing this 

[[x[0], x[1] + 's', x[2] + '1'] for x in a]

...unsurprisingly results in an error:

IndexError: list index out of range

Any ideas how to do this?

Upvotes: 1

Views: 91

Answers (4)

Mykola Zotko
Mykola Zotko

Reputation: 17804

You can use the function zip():

a = [['red', 'square', 'up'],
     ['red', 'circle'],
     ['blue', 'triangle'],
     ['blue', 'square'],
     ['blue', 'octagon'],
     ['blue', 'diamond', 'down']]

endings = ['', 's', '1']

[[i + j for i, j in zip(i, endings)] for i in a]

Result:

[['red', 'squares', 'up1'],
 ['red', 'circles'],
 ['blue', 'triangles'],
 ['blue', 'squares'],
 ['blue', 'octagons'],
 ['blue', 'diamonds', 'down1']]

Upvotes: 3

DirtyBit
DirtyBit

Reputation: 16772

Get the len() of each elem in the list, and manipulate accordingly:

For understanding:

a = [['red', 'square', 'up'], ['red', 'circle'],
     ['blue', 'triangle'], ['blue', 'square'],
     ['blue', 'octagon'], ['blue', 'diamond', 'down']]

for elem in a:
    if len(elem) < 3:
        print([elem[0], elem[1] + "s"])
    else:
        print([elem[0], elem[1] + "s " + elem[2]+ "1"])

OUTPUT:

['red', 'squares up1']
['red', 'circles']
['blue', 'triangles']
['blue', 'squares']
['blue', 'octagons']
['blue', 'diamonds down1']

Using list-comprehension:

print([[x[0], x[1] + 's'] if len(x) < 3 else [x[0], x[1] + 's', x[2] + '1'] for x in a])

Upvotes: 1

The Weckness
The Weckness

Reputation: 111

This worked for me in case you only want to check whether the nested list has a len > 2:

[[x[0], x[1] + 's'] if len(x) == 2 else [x[0], x[1]+'s', x[2]+'1'] for x in a]

Output:

[['red', 'squares', 'up1'],
 ['red', 'circles'],
 ['blue', 'triangles'],
 ['blue', 'squares'],
 ['blue', 'octagons'],
 ['blue', 'diamonds', 'down1']]

Upvotes: 0

ingvar
ingvar

Reputation: 4377

You got error because sometimes your x has 2 elements and sometimes 3. I modified your code to handle this:

a = [['red', 'square', 'up'], ['red', 'circle'], 
     ['blue', 'triangle'], ['blue', 'square'], 
     ['blue', 'octagon'], ['blue', 'diamond', 'down']]
t = [[x[0], x[1] + 's'] if len(x) == 2 else [x[0], x[1] + 's', x[2] + '1'] for x in a]
print(t)

Output:

[['red', 'squares', 'up1'], ['red', 'circles'], ['blue', 'triangles'], ['blue', 'squares'], ['blue', 'octagons'], ['blue', 'diamonds', 'down1']]

Upvotes: 2

Related Questions