python : generate sublists depending on the elements values

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.

Solution

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

[[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)