通过自定义键 function 将列表拆分为列表字典

Question

I implemented a small helper function that splits a list into a dictionary by a custom key function. The following implementation works, but feels rather un-pythonic:

def classify_list_by_key(l, key_func):
    result = defaultdict(list)

    for item in l:
        result[key_func(item)].append(item)

    return result

I tried using the following dict comprehension, which obviously has horrible performance with a lot of unnecessary calls to key_func and quadratic run-time.

def classify_list_by_key_dict_comprehension(l, key_func):
    return {key_func(x): [y for y in objects if key_func(x) == key_func(y)] for x in l}

---

Sample in- and output:

str_list = "these are some test strings for testing classify_list_by_key function".split()
print(classify_list_by_key(str_list, len))

produces

defaultdict(<class 'list'>, {5: ['these'], 3: ['are', 'for'], 4: ['some', 'test'], 7: ['strings', 'testing'], 20: ['classify_list_by_key'], 8: ['function']})

---

Is there a better / more pythonic way to achieve this using built-ins or standard library modules?

Answer 1

It is possible using groupby from itertools , but it requires a sorted list for that purpose.

Here is the code for that:

from itertools import groupby
str_list = "these are some test strings for testing classify_list_by_key function".split()
str_list.sort(key=len)
res = {k: list(g) for k, g in groupby(str_list, key=len)}
print(res)

Output

{3: ['are', 'for'], 4: ['some', 'test'], 5: ['these'], 7: ['strings', 'testing'], 8: ['function'], 20: ['classify_list_by_key']}

Answer 2

The following is an alternative implementation which does not force to use a defaultdict:

def classify_list_by_key(l, key_func):
    result = {}

    for item in l:
        result.setdefault(key_func(item), []).append(item)

    return result

str_list = "these are some test strings for testing classify_list_by_key function".split()
print(classify_list_by_key(str_list, len))