Python推导式构建集合字典

问题描述

简化以下代码，算法复杂度为 O(N)

from collections import defaultdict

data = [('math', 'a'), ('math', 'a'), ('math', 'b'), ('english', 'c')]
result = defaultdict(set)
for i in data:
    subject, student = i
    result[subject].add(student)
print(result)
# {'math': {'a', 'b'}, 'english': {'c'}}
1
2
3
4
5
6
7
8
9

解决方案一：列表循环调用 O(N)

from collections import defaultdict

data = [('math', 'a'), ('math', 'a'), ('math', 'b'), ('english', 'c')]
result = defaultdict(set)
_ = {result[k].add(v) for k, v in data}
print(result)
# {'math': {'a', 'b'}, 'english': {'c'}}
1
2
3
4
5
6
7

解决方案二：一行式 O(N²)

data = [('math', 'a'), ('math', 'a'), ('math', 'b'), ('english', 'c')]
result = {key: {value for _key, value in data if _key == key} for key, _ in data}
print(result)
1
2
3

解决方案三：groupby O(NlogN)

from itertools import groupby
from operator import itemgetter

data = [('math', 'a'), ('math', 'a'), ('math', 'b'), ('english', 'c')]
result = {k: {x[1] for x in g} for k, g in groupby(data, key=itemgetter(0))}
print(result)
# {'math': {'a', 'b'}, 'english': {'c'}}
1
2
3
4
5
6
7

解决方案四：reduce O(N)

from functools import reduce
from collections import defaultdict

data = [('math', 'a'), ('math', 'a'), ('math', 'b'), ('english', 'c')]
result = reduce(lambda d, kv: d[kv[0]].add(kv[1]) or d, data, defaultdict(set))
print(result)
# {'math': {'a', 'b'}, 'english': {'c'}}
1
2
3
4
5
6
7

参考文献

1. collections — Python 文档
2. Python get a set map with list comprehension in one line
3. One-liner to create dictionary of lists
4. One-step initialization of defaultdict that appends to list?