一文详解Python如何优雅地对数据进行分组

假设我们有这样一种数据：

data = [
    (\"apple\", 30), (\"apple\", 35),
    (\"apple\", 32), (\"pear\", 60),
    (\"pear\", 32), (\"pear\", 60),
    (\"banana\", 102), (\"banana\", 104)
]

# 我们希望变成如下格式
\"\"\"
[(\'apple\', [30, 35, 32]), 
 (\'pear\', [60, 32, 60]), 
 (\'banana\', [102, 104])]
\"\"\"

如果是你的话，你会怎么做呢？很容易想到的一种解决方案是构造一个字典：

data = [
    (\"apple\", 30), (\"apple\", 35),
    (\"apple\", 32), (\"pear\", 60),
    (\"pear\", 32), (\"pear\", 60),
    (\"banana\", 102), (\"banana\", 104)
]

data_dict = {}
for name, count in data:
    if name not in data_dict:
        data_dict[name] = []
    data_dict[name].append(count)
print(data_dict)
\"\"\"
{\'apple\': [30, 35, 32], 
 \'pear\': [60, 32, 60], 
 \'banana\': [102, 104]}
\"\"\"
print(list(data_dict.items()))
\"\"\"
[(\'apple\', [30, 35, 32]), 
 (\'pear\', [60, 32, 60]), 
 (\'banana\', [102, 104])]
\"\"\"

这种方案完全没有问题，不过我们还可以写的更优雅一些，也就是使用字典的 setdefault 方法：

data = [
    (\"apple\", 30), (\"apple\", 35),
    (\"apple\", 32), (\"pear\", 60),
    (\"pear\", 32), (\"pear\", 60),
    (\"banana\", 102), (\"banana\", 104)
]

data_dict = {}
for name, count in data:
    # setdefault(k, v) 含义如下
    # 当 k 不存在时，将 k: v 设置在字典中，并返回 v
    # 当 k 存在时，直接返回 k 对应值
    data_dict.setdefault(name, []).append(count)

print(list(data_dict.items()))
\"\"\"
[(\'apple\', [30, 35, 32]), 
 (\'pear\', [60, 32, 60]), 
 (\'banana\', [102, 104])]
\"\"\"

setdefault 是一个非常方便的方法，但是使用频率却不怎么高，或者说该方法不太让人喜欢。主要是每次调用都要给一个初始值，比如代码中的空列表 []。另外这里的初始值可以任意，如果你希望添加的时候还能实现去重效果，那么就将空列表换成空集合即可。

或者我们还可以使用 defaultdict，它位于 collections 模块中。

from collections import defaultdict

data = [
    (\"apple\", 30), (\"apple\", 35),
    (\"apple\", 32), (\"pear\", 60),
    (\"pear\", 32), (\"pear\", 60),
    (\"banana\", 102), (\"banana\", 104)
]

# 里面接收一个 callable
# 当访问的 k 不存在时，返回 callable 调用之后的值
data_dict1 = defaultdict(list)
for name, count in data:
    data_dict1[name].append(count)

print(list(data_dict1.items()))
\"\"\"
[(\'apple\', [30, 35, 32]), 
 (\'pear\', [60, 32, 60]), 
 (\'banana\', [102, 104])]
\"\"\"

# 也可以指定为 set
data_dict2 = defaultdict(set)
for name, count in data:
    data_dict2[name].add(count)

print(list(data_dict2.items()))
\"\"\"
[(\'apple\', {32, 35, 30}), 
 (\'pear\', {32, 60}), 
 (\'banana\', {104, 102})]
\"\"\"

总的来说，defaultdict 和字典的 setdefault 方法非常类似，我们使用 setdefault 即可。

当然啦，关于分组，还有一种特殊情况，就是词频统计。假设我们想统计可迭代对象中，每个元素出现的次数该怎么做呢？

data = [\"apple\", \"apple\", \"apple\",
        \"pear\", \"pear\", \"pear\",
        \"banana\", \"banana\"]

data_dict = {}
for item in data:
    # 此处不能使用 setdefault，因为它是函数
    # .setdefault(item, 0) += 1 是不符合语法规则的
    if item not in data_dict:
        data_dict[item] = 0
    data_dict[item] += 1

print(data_dict)
\"\"\"
{\'apple\': 3, \'pear\': 3, \'banana\': 2}
\"\"\"

# 或者使用 defaultdict
from collections import defaultdict
data_dict = defaultdict(int)
for item in data:
    data_dict[item] += 1
print(data_dict)
\"\"\"
defaultdict(<class \'int\'>, 
            {\'apple\': 3, \'pear\': 3, \'banana\': 2})
\"\"\"

然而说到词频统计，我们还可以使用 collections 下的 Counter 类。

from collections import Counter

data = [\"apple\", \"apple\", \"apple\",
        \"pear\", \"pear\", \"pear\",
        \"banana\", \"banana\"]

data_dict = Counter(data)
# 直接搞定，Counter 已经包含了我们之前的逻辑
print(data_dict)
\"\"\"
Counter({\'apple\': 3, \'pear\': 3, \'banana\': 2})
\"\"\"
# Counter 继承 dict，除了支持字典操作之外
# 还提供了很多其它操作，其中一个就是 most_common
# 用于选择出现频率最高的几个元素
print(data_dict.most_common(2))
\"\"\"
[(\'apple\', 3), (\'pear\', 3)]
\"\"\"

还是很简单的。