Elio Maisonneuve
Elio Maisonneuve

Reputation: 369

How to do intermediate process in list comprehension

I have a list of string that look like that :

['1 2 3 4 5', '1 2 3 4 5',...]

And I want for example only the third and fifth element of each string :

[(3,5), (3,5),...]

how can I do something like this :

[(ans[2],ans[4]); ans.split() for ans in answers]

?

Upvotes: 3

Views: 131

Answers (2)

Moses Koledoye
Moses Koledoye

Reputation: 78546

Use operator.itemgetter to get multiple items from different indices:

from operator import itemgetter

lst = ['1 2 3 4 5', '1 2 3 4 5']

f = itemgetter(2, 4)             # get third and fifth items
r = [f(x.split()) for x in lst]
print(r)
# [('3', '5'), ('3', '5')]

Cool thing is, it also works with slice objects:

f = itemgetter(2, 4, slice(3, 0, -1))
r = [f(x.split()) for x in lst]
print(r)
# [('3', '5', ['4', '3', '2']), ('3', '5', ['4', '3', '2'])]

Upvotes: 9

willeM_ Van Onsem
willeM_ Van Onsem

Reputation: 476584

You first perform a mapping, for instance using a genrator:

(ans.split() for ans in answers)

and then you iterate over the generator, and process the intermediate result further, like:

[(imm[2],imm[4]) for imm in (ans.split() for ans in answers)]

generating:

>>> [(imm[2],imm[4]) for imm in (ans.split() for ans in answers)]
[('3', '5'), ('3', '5')]

Upvotes: 6

Related Questions