itertools.groupby()分组不正确

Question

我有这些数据：

self.data = [(1, 1, 5.0),
             (1, 2, 3.0),
             (1, 3, 4.0),
             (2, 1, 4.0),
             (2, 2, 2.0)]

当我运行这段代码时：

for mid, group in itertools.groupby(self.data, key=operator.itemgetter(0)):

对于 list(group)，我得到：

[(1, 1, 5.0),
 (1, 2, 3.0),
 (1, 3, 4.0)]

这正是我想要的结果。

但是如果我用1代替0

for mid, group in itertools.groupby(self.data, key=operator.itemgetter(1)):

来根据元组中的第二个数字进行分组，我只得到：

[(1, 1, 5.0)]

尽管还有其他元组在那个位置上有“1”。

Answer 1

下面这个内容解决了Python中itertools.groupby的一些烦人问题。

def groupby2(l, key=lambda x:x, val=lambda x:x, agg=lambda x:x, sort=True):
    if sort:
        l = sorted(l, key=key)
    return ((k, agg((val(x) for x in v))) \
        for k,v in itertools.groupby(l, key=key))

具体来说，

1. 它不要求你必须先对数据进行排序。
2. 你可以不必只用key作为命名参数。
3. 输出结果是一个干净的生成器，格式是tuple(key, grouped_values)，其中的值是通过第三个参数指定的。
4. 可以很方便地应用聚合函数，比如求和或平均值。

示例用法

import itertools
from operator import itemgetter
from statistics import *

t = [('a',1), ('b',2), ('a',3)]
for k,v in groupby2(t, itemgetter(0), itemgetter(1), sum):
  print(k, v)

这段代码会打印出，

a 4
b 2

可以试试这段代码

Answer 2

不需要排序的变体（通过字典实现）。这样做在性能上应该会更好。

def full_group_by(l, key=lambda x: x):
    d = defaultdict(list)
    for item in l:
        d[key(item)].append(item)
    return d.items()

Answer 3

itertools.groupby 是一个可以把相邻的、具有相同特征的项目聚集在一起的工具。如果你想把所有具有相同特征的项目都找出来，首先需要对 self.data 进行排序。

for mid, group in itertools.groupby(
    sorted(self.data,key=operator.itemgetter(1)), key=operator.itemgetter(1)):