Python Comprehension With Range

Question

>>> x = {'a':2, 'b':3}
>>> [key for key in x for i in range(x[key])]
['b', 'b', 'b', 'a', 'a']

Is there a better way to write this code without the for i in range(x[key]). Or is there a better way to write this?

Answer 1

or use a Counter:

from collections import Counter

x = Counter({'a':2, 'b':3})
print (list(x.elements()))
# ['b', 'b', 'b', 'a', 'a']

Answer 2

You could use the itertools module to do the work:

from itertools import chain, starmap, repeat

list(chain.from_iterable(starmap(repeat, x.items())))

This is how the Counter.elements() method is implemented, which you could use directly:

from collections import Counter

list(Counter(x).elements())

If you want to stick to a list comprehension, then at least use the itertools.repeat() function still, and use dict.items() (or even dict.iteritems() in Python 2) to give you key and associated value pairs:

from itertools import repeat

[c for key, count in x.items() for c in repeat(key, count)]

All three produce the same output:

>>> from itertools import chain, starmap, repeat
>>> from collections import Counter
>>> x = {'a': 2, 'b': 3}
>>> list(chain.from_iterable(starmap(repeat, x.items())))
['a', 'a', 'b', 'b', 'b']
>>> list(Counter(x).elements())
['a', 'a', 'b', 'b', 'b']
>>> [c for key, count in x.items() for c in repeat(key, count)]
['a', 'a', 'b', 'b', 'b']

where the relative order of the keys is dependent on the dictionary insertion and deletion order.