vonBulow
vonBulow

Reputation: 87

List comprehension with some kind of loop interdependence?

I'm sure plenty of questions have been asked about list comprehensions but I have one about a very particular case involving nested loops with co-reference that I can't find answered elsewhere.

Say you have a dictionary that satisfies two constraints: (1) every value in the dictionary is a list of some arbitrary length and (2) there exists some list outside the dictionary whose elements are a subset of the keys in the dictionary (this is necessary for reasons beyond the scope of this question).

Got that? It looks a little like this:

someDict = {'a':[0, 1], 'b':[2, 3], 'c':[4, 5], 'd':[6, 7], 'e':[8, 9]}
supraList = ['b', 'c', 'd']

Now what I want to do is:

1) cycle through only those lists of someDict whose keys are also elements of supraList;

2) add the elements of those someDict lists (not the lists themselves) to some new supra-list.

The effect is something like this:

... CODE ...
newSupraList = [2, 3, 4, 5, 6, 7]

i.e. supraList contained 'b', 'c' and 'd', so newSupraList gets the combined elements of the lists associated with 'b', 'c' and 'd' in someDict.

Achieving this with a loop goes a little like:

newSupraList = []
for i in supraList:
    for j in someDict[i]:
        newSupraList.append(j)

Can this be achieved with a comprehension? The best comprehension I could write (and it was ugly) came out with:

newSupraList = [[2, 3], [4, 5], [6, 7]]

I'm sure there's a really simple way of iterating over a list such as that to make it flat despite all of the above, but I thought I'd get some fresh eyes to take a stab at the problem for some new angles. The loops may be better for readability, but I think it's fun in any case!

Upvotes: 2

Views: 58

Answers (2)

zondo
zondo

Reputation: 20336

Yes, sir:

newSupraList = [item for key in supraList for item in someDict[key]]

Upvotes: 4

timgeb
timgeb

Reputation: 78690

Sure, you can have two for loops in a comprehension as well.

>>> [x for y in supraList for x in someDict[y]]
[2, 3, 4, 5, 6, 7]

Alternatively, you can use chain from itertools.

With unpacking:

>>> from itertools import chain
>>> list(chain(*(someDict[x] for x in supraList)))
[2, 3, 4, 5, 6, 7]

With chain.from_iterable:

>>> list(chain.from_iterable(someDict[x] for x in supraList))
[2, 3, 4, 5, 6, 7]

Upvotes: 5

Related Questions