python : generate sublists depending on the elements values

Question

I have a list of elements, and I want to generate an arbitrary number of sub-lists from that list. All the elements of a given sub-list must have the same output from a function f.

For example :

>>> mylist = [1,2,3,4,5,6] # my list to split into sublist
>>> sublists = list() # my sublists
>>> def f(e) : # the function used to discriminate the elements
...  return e%2
... 
>>> for ret in [0,1] :
...   sublists.append([e for e in mylist if f(e) == ret])
... 
>>> print(sublists)
[[2, 4, 6], [1, 3, 5]]

This works great, but only because I know exactly which values f can return: 0, or 1. So I can iterate on these output values and then create the sub-lists.

But what if I don't know beforehand the values that f can return? Is there a good Pythonic way to do it? I would like to avoid having to iterate over mylist just to see what f can output on all the elements.

Answer 1

You could use a collections.defaultdict:

from collections import defaultdict

out = defaultdict(list)
for e in mylist:
    out[f(e)].append(e)

Now as long as f(e) is hashable for all appropriate values of e, you will get an output like

defaultdict(<class 'list'>, {0: [2, 4, 6], 1: [1, 3, 5]})

If you want a list of these sub-lists, you can then do:

sublists = list(out.values())

However note that dictionaries aren't guaranteed to retain a sensible order; each sublist will remain ordered, but you could get

[[2, 4, 6], [1, 3, 5]]

or

[[1, 3, 5], [2, 4, 6]]

in this example.

---

The long-form equivalent to show what defaultdict does:

out = {}
for e in mylist:
    key = f(e)
    if key not in out:
        out[key] = []
    out[key].append(e)

Answer 2

Just for diversity:

from itertools import groupby
from operator import itemgetter

mylist = [1, 2, 3, 4, 5, 6] 
f = lambda x: x % 2

zipt = sorted((f(x), x) for x in mylist)

for x, y in groupby(zipt, key=itemgetter(0)):
    print(x, [z[1] for z in y])

prints

0 [2, 4, 6]
1 [1, 3, 5]