List comprehension from dict of lists

Question

I have a dict-

a = {'b': [1,2,3], 'c':[4,5,6]}

I want to use list comprehension only to achieve this output-

[['c', 4], ['c', 5], ['c', 6], ['b', 1], ['b', 2], ['b', 3]]

A simple for loop gets it done with -

x = []
for k, v in a.iteritems():
    for i in v:
        x.append([k, i])

Tried to convert it to list comprehension, I did this-

[[k,i] for i in v for k, v in a.items()]

But weirdly for me, I got an output

[['c', 1], ['b', 1], ['c', 2], ['b', 2], ['c', 3], ['b', 3]]

What should be the right list comprehension and why is my list comprehension not working?

Mohit Solanki · Answer 1 · 2018-08-28T07:32:47.240

9

b = [[i, j] for i in a for j in a[i]]

for nested for loops in list comprehension the first loop will be the one whose variable you will be using in the second one, like here for example i is used in the second loop, nested for loops in list comprehensions are hard to read therefore you should better avoid it.

edited Aug 28 '18 at 07:32

answered Aug 28 '18 at 07:11

Mohit Solanki

2,122
12
20

Why on right of `a.items()` and not on the left in the list comprehension? – Sushant Aug 28 '18 at 07:19
sorry, didn't get your question? – Mohit Solanki Aug 28 '18 at 07:23
1

How do we decide order of the for loop in list comprehension? – Sushant Aug 28 '18 at 07:24
for nested `for` loops in list comprehension the first loop will be the one whose variable you will be using in the second one, like here for example `i` is used in the second loop, nested for loops in list comprehensions are hard to read therefore you should better avoid it. – Mohit Solanki Aug 28 '18 at 07:28
Thanks, upvoted your answer, a complete answer would have been this explanation in it as well – Sushant Aug 28 '18 at 07:30

madjaoue · Answer 2 · 2018-08-28T07:30:04.493

4

You should get k,v first, then iterate over v :

a = {'b': [1,2,3], 'c':[4,5,6]}
print([[k,i] for k, v in a.items() for i in v])

output :

[['c', 4], ['c', 5], ['c', 6], ['b', 1], ['b', 2], ['b', 3]]

Note :

In [[k,i] for i in v for k, v in a.items()], v is not defined when you try to iterate over it.

@Skam has a great example for this : how to interpret double for loop comprehension

# Without list comprehension
list_of_words = []
for sentence in text:
    for word in sentence:
       list_of_words.append(word)
return list_of_words

is equivalent to :

[word for sentence in text for word in sentence]

edited Aug 28 '18 at 07:30

answered Aug 28 '18 at 07:12

madjaoue

5,104
2
19
31

Why on right of `a.items()` and not on the left in the list comprehension? – Sushant Aug 28 '18 at 07:18
It has to do with [How to interpret double for loop comprehension](https://stackoverflow.com/a/36734643/9057497). In your implementation, v is not yet defined when you try to iterate over it. – madjaoue Aug 28 '18 at 07:24
It'll be helpful if you add this in your answer as well as that'll completely answer my question – Sushant Aug 28 '18 at 07:26
1

Thanks, upvoted your answer, argo added a complete answer before you did so accepted his answer – Sushant Aug 28 '18 at 07:31

score 4 · Accepted Answer · answered Aug 28 '18 at 07:23

You were almost there. The main problem you were facing is due to the order of the for loop.

The order of for loop inside the list comprehension is based on the order in which they appear in traditional loop approach. Outer most loop comes first, and then the inner loops subsequently.

a = {'b': [1,2,3], 'c':[4,5,6]}
x = []
for k, v in a.items():
    for i in v:
        x.append([k, i])

print(x)

print([[k,i] for i in v for k, v in a.items()])
print([[k,i] for k, v in a.items() for i in v])

OUTPUT

[['b', 1], ['b', 2], ['b', 3], ['c', 4], ['c', 5], ['c', 6]]
[['b', 4], ['c', 4], ['b', 5], ['c', 5], ['b', 6], ['c', 6]]
[['b', 1], ['b', 2], ['b', 3], ['c', 4], ['c', 5], ['c', 6]]

Brown Bear · Answer 4 · 2018-08-28T07:39:26.607

you can try to use itertools.product

from itertools import product, chain
a = {'b': [1,2,3], 'c':[4,5,6]}
list(chain(*[product(k, v) for k, v in a.items()]))

result is

[('b', 1), ('b', 2), ('b', 3), ('c', 4), ('c', 5), ('c', 6)]

if you strongly need the list of lists, you can do

list(chain(*[[list(item) for item in product(k, v)] for k, v in a.items()]))

the output is:

[['b', 1], ['b', 2], ['b', 3], ['c', 4], ['c', 5], ['c', 6]]

and some tests of the perfomance

In [6]: %timeit [[i, j] for i in a for j in a[i]]
618 ns ± 5.34 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

In [7]: %timeit list(chain(*[product(k, v) for k, v in a.items()]))
1.26 µs ± 19.4 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

In [8]: %timeit list(chain(*[[list(item) for item in product(k, v)] for k, v in a.items()]))
2.61 µs ± 49.9 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

List comprehension from dict of lists

4 Answers4