CODEWITHSUNDEEP
pythonlist-comprehension
Sushant
Sushant

Reputation: 3669

List comprehension from dict of lists

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?

Upvotes: 3

Views: 5102

Answers (4)

user8060120
user8060120

Reputation:

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)

Upvotes: 1

Mohit Solanki
Mohit Solanki

Reputation: 2130

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.

Upvotes: 8

madjaoue
madjaoue

Reputation: 5224

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]

Upvotes: 4

Arghya Saha
Arghya Saha

Reputation: 5713

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]]

Upvotes: 3

Related Questions