Split/partition list based on invariant/hash?

Question

I have a list [a1,21,...] and would like to split it based on the value of a function f(a). For example if the input is the list [0,1,2,3,4] and the function def f(x): return x % 3, I would like to return a list [0,3], [1,4], [2], since the first group all takes values 0 under f, the 2nd group take value 1, etc...

Something like this works: return [[x for x in lst if f(x) == val] for val in set(map(f,lst))],

But it does not seem optimal (nor pythonic) since the inner loop unnecessarily scans the entire list and computes same f values of the elements several times. I'm looking for a solution that would compute the value of f ideally once for every element...

Tim Peters · Accepted Answer · 2018-08-28T00:30:47.957

3

If you're not irrationally ;-) set on a one-liner, it's straightforward:

from collections import defaultdict

lst = [0,1,2,3,4]
f = lambda x: x % 3

d = defaultdict(list)
for x in lst:
    d[f(x)].append(x)
print(list(d.values()))

displays what you want. f() is executed len(lst) times, which can't be beat

EDIT: or, if you must:

from itertools import groupby
print([[pair[1] for pair in grp]
       for ignore, grp in
       groupby(sorted((f(x), x) for x in lst),
               key=lambda pair: pair[0])])

That doesn't require that f() produce values usable as dict keys, but incurs the extra expense of a sort, and is close to incomprehensible. Clarity is much more Pythonic than striving for one-liners.

edited Aug 28 '18 at 00:30

answered Aug 27 '18 at 23:55

Tim Peters

67,464
13
126
132

2

Would someone care to explain their downvote? I'd like to hear it. It answers the OP's question, in two different ways that do what they asked for. – Tim Peters Aug 28 '18 at 05:20
Cool, thank you. I like the 1st solution. Was not aware of the defaultdict class, very useful! – Jake B. Aug 28 '18 at 14:38
Yup, `defaultdict` can be handy. With a regular dict, you could test for whether the key already exists first, or do `d.setdefault(f(x), []).append(x)` without testing first. – Tim Peters Aug 28 '18 at 15:32

pylang · Answer 2 · 2018-08-29T00:21:21.207

0

@Tim Peters is right, and here is a mentioned setdefault and another itertool.groupby option.

Given

import itertools as it


iterable = range(5)
keyfunc = lambda x: x % 3

Code

setdefault

d = {}
for x in iterable:
    d.setdefault(keyfunc(x), []).append(x)

list(d.values())

groupby

[list(g) for _, g in it.groupby(sorted(iterable, key=keyfunc), key=keyfunc)]

Split/partition list based on invariant/hash?

2 Answers2