Making a list comprehension for a dict within a list within a list

Question

I need to create a list comprehension that extracts values from a dict within a list within a list, and my attempts so far are failing me. The object looks like this:

MyList=[[{'animal':'A','color':'blue'},{'animal':'B','color':'red'}],[{'animal':'C','color':'blue'},{'animal':'D','color':'Y'}]] 

I want to extract the values for each element in the dict/list/list so that I get two new lists:

Animals=[[A,B],[C,D]]
Colors=[[blue,red],[blue,Y]]

Any suggestions? Doesn't necessarily need to use a list comprehension; that's just been my starting point so far. Thanks!

Answer 1

Animals = [[d['animal'] for d in sub] for sub in MyList]
Colors = [[d['color'] for d in sub] for sub in MyList]

Gives the desired result:

[['A', 'B'], ['C', 'D']]
[['blue', 'red'], ['blue', 'Y']]  # No second 'red'.

---

What I have done here is take each sub-list, then each dictionary, and then access the correct key.

Answer 2

In a single assignment (with a single list comprehension, and the help of map and zip):

Colors, Animals =  map(list, 
                       zip(*[map(list, 
                                 zip(*[(d['color'], d['animal']) for d in a])) 
                             for a in MyList]))

If you are fine with tuples, you can avoid the two calls to map => list

EDIT:

Let's see it in some details, by decomposing the nested comprehension. Let's also assume MyList have m elements, for a total of n objects (dictionaries).

[[d for d in sub] for sub in MyList]

This would iterate through every dictionary in the sublists. For each of them, we create a couple with its color property in the first element and its animal property in the second one:

(d['color'], d['animal'])

So far, this will take time proportional to O(n) - exatly n elements will be processed.

print [[(d['color'], d['animal']) for d in sub] for sub in MyList]

Now, for each of the m sublists of the original list, we have one list of couples that we need to unzip, i.e. transform it into two lists of singletons. In Python, unzip is performed using the zip function by passing a variable number of tuples as arguments (the arity of the first tuple determines the number of tuples it outputs). For instance, passing 3 couples, we get two lists of 3 elements each

>>> zip((1,2), (3,4), (5,6))  #Prints [(1, 3, 5), (2, 4, 6)]

To apply this to our case, we need to pass array of couples to zip as a variable number of arguments: that's done using the splat operator, i.e. *

[zip(*[(d['color'], d['animal']) for d in sub]) for sub in MyList]

This operation requires going through each sublist once, and in turn through each one of the couples we created in the previous step. Total running time is therefore O(n + n + m) = O(n), with approximatively 2*n + 2*m operations.

So far we have m sublists, each one containing two tuples (the first one will gather all the colors for the sublist, the second one all the animals). To obtain two lists with m tuples each, we apply unzip again

zip(*[zip(*[(d['color'], d['animal']) for d in sub]) for sub in MyList]

This will require an additional m steps - the running time will therefore stay O(n), with approximatively 2*n + 4*m operations.

For sake of simplicity we left out mapping tuples to lists in this analysis - which is fine if you are ok with tuples instead.

Tuples are immutable, however, so it might not be the case. If you need lists of lists, you need to apply the list function to each tuple: once for each of the m sublists (with a total of 2*n elements), and once for each of the 2 first level lists, i.e. Animals and Colors, (which have a total of m elements each). Assuming list requires time proportional to the length of the sequence it is applied to, this extra step requires 2*n + 2*m operations, which is still O(n).