Python itertools.groupby() 邻近分组

说明

《Python 教程》持续更新中，提供建议、纠错、催更等加作者微信: gairuo123（备注：pandas教程）和关注公众号「盖若」ID: gairuo。跟作者学习，请进入 Python学习课程。欢迎关注作者出版的书籍：《深入浅出Pandas》和《Python之光》。

itertools.groupby() 是 Python itertools 模块中的一个函数，用于根据指定的 key 函数对可迭代对象中的元素进行分组。它返回一个迭代器，每次迭代会产生一个键和一个对应的迭代器，迭代器生成与当前键相等的元素。

语法

语法为 itertools.groupby(iterable, key=None) 创建一个迭代器，返回 iterable 中连续的键和组。key 是一个计算元素键值函数。如果未指定或为 None，key 缺省为恒等函数（identity function），返回元素不变。一般来说，iterable 需用同一个键值函数预先排序。

参数和返回值：

iterable：要进行分组的可迭代对象（如列表、元组、字符串等）。
key：用于分组的函数，如果未提供，则默认为 None，表示直接使用元素自身作为键。
返回值：返回一个迭代器，每次迭代产生一个 (key, group) 对，其中 key 是分组的键，group 是与该键对应的元素迭代器。

注意事项：

输入的可迭代对象需要是按照分组键进行排序的，否则结果可能不符合预期。
itertools.groupby() 只能识别连续相同的元素，并将它们作为一组处理。

这种行为与 SQL 的 GROUP BY 操作不同，SQL 的操作会忽略输入的顺序将相同键值的元素分在同组中。itertools.groupby 将邻近的相同元素分组。

import itertools

foo = itertools.groupby([0,0,2,2,3,3,2,2])
[*foo]
'''
[(0, <itertools._grouper at 0x7f9d31bf7910>),
 (2, <itertools._grouper at 0x7f9d31bf4820>),
 (3, <itertools._grouper at 0x7f9d31bf7fa0>),
 (2, <itertools._grouper at 0x7f9d31bf7c70>)]
'''

由于 itertools.groupby() 返回的组本身也是一个迭代器，它与 groupby() 共享底层的可迭代对象。因为源是共享的，当 groupby() 对象向后迭代时，前一个组将消失。因此如果稍后还需要使用的话，可保存为列表：

import itertools

foo = itertools.groupby([0,0,2,2,3,3,2,2], key=lambda x: str(x+1))

groupskeys, groups = [], []
for gk, g in foo:
    groupskeys.append(gk)
    groups.append([*g])

groupskeys
# ['1', '3', '4', '3']
groups
# [[0, 0], [2, 2], [3, 3], [2, 2]]

groupby() 大致相当于：

class groupby:
    # [k for k, g in groupby('AAAABBBCCDAABBB')] --> A B C D A B
    # [list(g) for k, g in groupby('AAAABBBCCD')] --> AAAA BBB CC D
    def __init__(self, iterable, key=None):
        if key is None:
            key = lambda x: x
        self.keyfunc = key
        self.it = iter(iterable)
        self.tgtkey = self.currkey = self.currvalue = object()
    def __iter__(self):
        return self
    def __next__(self):
        self.id = object()
        while self.currkey == self.tgtkey:
            self.currvalue = next(self.it)    # Exit on StopIteration
            self.currkey = self.keyfunc(self.currvalue)
        self.tgtkey = self.currkey
        return (self.currkey, self._grouper(self.tgtkey, self.id))
    def _grouper(self, tgtkey, id):
        while self.id is id and self.currkey == tgtkey:
            yield self.currvalue
            try:
                self.currvalue = next(self.it)
            except StopIteration:
                return
            self.currkey = self.keyfunc(self.currvalue)

示例

基本用法

import itertools

data = [('a', 1), ('a', 2), ('b', 3), ('b', 4), ('a', 5)]

# 使用 itertools.groupby 对元组列表按照第一个元素进行分组
groups = itertools.groupby(data, key=lambda x: x[0])

for key, group in groups:
    print(key, list(group))

# 输出：
'''
a [('a', 1), ('a', 2)]
b [('b', 3), ('b', 4)]
a [('a', 5)]
'''

使用字符串分组

import itertools

data = ['apple', 'bat', 'bar', 'atom', 'book']

# 使用 itertools.groupby 对字符串列表按照首字母进行分组
groups = itertools.groupby(data, key=lambda x: x[0])

for key, group in groups:
    print(key, list(group))

# 输出：
'''
a ['apple', 'atom']
b ['bat', 'bar', 'book']
'''

不使用 key 函数

import itertools

data = [1, 2, 2, 3, 3, 3]

# 使用 itertools.groupby 对整数列表按照元素自身进行分组
groups = itertools.groupby(data)

for key, group in groups:
    print(key, list(group))

# 输出：
'''
1 [1]
2 [2, 2]
3 [3, 3, 3]
'''

应用场景

数据分析：在处理数据集时，根据某些特征进行分组统计。
日志分析：按照时间或者日志级别分组分析日志信息。
文本处理：按照字母、单词或者段落对文本进行分组处理。

实用案例：分组统计成绩

假设我们有一组学生的姓名和成绩，我们希望按照姓名对他们的成绩进行分组统计。

示例数据：

data = [
    ('Alice', 80),
    ('Bob', 85),
    ('Alice', 90),
    ('Bob', 75),
    ('Alice', 95),
]

目标：

按照姓名分组，计算每个学生的总成绩和平均成绩：

{
    'Alice': {'total': 265, 'average': 88.33},
    'Bob': {'total': 160, 'average': 80.0}
}

实现代码：

import itertools
from collections import defaultdict

data = [
    ('Alice', 80),
    ('Bob', 85),
    ('Alice', 90),
    ('Bob', 75),
    ('Alice', 95),
]

# 使用 defaultdict 初始化一个空字典
results = defaultdict(lambda: {'total': 0, 'count': 0})

# 使用 itertools.groupby 按姓名分组并计算总成绩和次数
for key, group in itertools.groupby(data, key=lambda x: x[0]):
    total = sum(score for _, score in group)
    results[key]['total'] += total
    results[key]['count'] += len(list(group))

# 计算平均成绩
for name, stats in results.items():
    stats['average'] = stats['total'] / stats['count']

# 输出结果
print(dict(results))

解释：

初始化结果字典：使用 defaultdict 初始化一个空字典 results，用于存储每个学生的总成绩和平均成绩。
使用 itertools.groupby 进行分组：itertools.groupby(data, key=lambda x: x[0]) 按照姓名分组，并生成分组的迭代器。
计算每个学生的总成绩和次数：遍历分组后的迭代器，计算每个学生的总成绩和次数。
计算平均成绩：根据总成绩和次数计算每个学生的平均成绩。
输出结果：打印最终的分组统计结果字典。

输出结果：

{
    'Alice': {'total': 265, 'average': 88.33},
    'Bob': {'total': 160, 'average': 80.0}
}

这个实用案例展示了如何使用 itertools.groupby 对数据进行分组，并在分组后进行进一步的统计和计算，适用于需要按照某些特征对数据进行分组和分析的场景。

参考

https://docs.python.org/zh-cn/3/library/itertools.html#itertools.groupby

< filterfalse() 按假值筛选 Python 标准库 itertools 为高效循环而创建迭代器的函数 islice() 按索引切片迭代 >

更新时间：2024-07-03 22:06:04 标签：python groupby 分组