Python 新手突破瓶颈指南：通过 itertools.groupby 进行数据分组

MegaQi

发布于 2024-08-20 21:24:23

920

发布于 2024-08-20 21:24:23

文章被收录于专栏：非典型性程序员

在数据处理和分析中，我们常常需要根据某些条件对数据进行分组。itertools.groupby() 能够帮助我们轻松实现这一功能。

工作机制

itertools.groupby() 会对输入的可迭代对象进行分组，返回一个迭代器，该迭代器生成连续键值对，其中每个键值对的键是分组依据，值是一个生成器，生成对应分组的元素。

class groupby(object):
    """
    make an iterator that returns consecutive keys and groups from the iterable
    
      iterable
        Elements to divide into groups according to the key function.
      key
        A function for computing the group category for each element.
        If the key function is not specified or is None, the element itself
        is used for grouping.
    """

iterable: 要分组的可迭代对象。
key: 一个函数，用于指定分组依据。如果未指定，则使用元素本身进行分组。

基本示例

以下是一个简单的示例，演示如何使用 itertools.groupby() 对列表进行分组：

import itertools

data = [1, 1, 2, 2, 3, 3, 3, 4, 4, 4, 4]
grouped = itertools.groupby(data)

for key, group in grouped:
    print(f'Key: {key}, Group: {list(group)}')

输出：

Key: 1, Group: [1, 1]
Key: 2, Group: [2, 2]
Key: 3, Group: [3, 3, 3]
Key: 4, Group: [4, 4, 4, 4]

在这个例子中，itertools.groupby() 将连续相同的元素分组。

从底层逻辑来看，itertools.groupby() 通过遍历输入的可迭代对象，并根据键函数返回的值进行分组。如果键函数未指定，则默认使用元素本身作为键。需要注意的是，itertools.groupby() 只会对相邻的元素进行分组，因此在使用之前通常需要对数据进行排序。

比如，上边的 data 是个非排序数组

import itertools

data = [1, 4, 2, 4, 3, 3, 3, 4, 4, 4, 4]
grouped = itertools.groupby(data)

for key, group in grouped:
    print(f'Key: {key}, Group: {list(group)}')

则输出为：

Key: 1, Group: [1]
Key: 4, Group: [4]
Key: 2, Group: [2]
Key: 4, Group: [4]
Key: 3, Group: [3, 3, 3]
Key: 4, Group: [4, 4, 4, 4]

应用场景

1. 对列表进行分组

在实际应用中，通常需要对列表进行分组，例如根据值的相等性对元素分组：

import itertools

data = ['apple', 'banana', 'apple', 'orange', 'banana', 'apple']
data.sort()
grouped = itertools.groupby(data)

for key, group in grouped:
    print(f'Key: {key}, Group: {list(group)}')

输出：

Key: apple, Group: ['apple', 'apple', 'apple']
Key: banana, Group: ['banana', 'banana']
Key: orange, Group: ['orange']

2. 按照对象属性分组

itertools.groupby() 还可以用于根据对象的属性进行分组。例如，有一个包含多个字典的列表，可以按某个键的值进行分组：

import itertools

data = [
    {'name': 'Alice', 'age': 30},
    {'name': 'Bob', 'age': 25},
    {'name': 'Charlie', 'age': 30},
    {'name': 'David', 'age': 25}
]

data.sort(key=lambda x: x['age'])
grouped = itertools.groupby(data, key=lambda x: x['age'])

for key, group in grouped:
    print(f'Age: {key}')
    for item in group:
        print(f'    {item}')

输出：

Age: 25
    {'name': 'Bob', 'age': 25}
    {'name': 'David', 'age': 25}
Age: 30
    {'name': 'Alice', 'age': 30}
    {'name': 'Charlie', 'age': 30}

3. 处理日志文件

在分析日志文件时，可以根据时间戳或日志级别对日志进行分组：

import itertools

logs = [
    '2023-01-01 INFO Start processing',
    '2023-01-01 ERROR Something went wrong',
    '2023-01-02 INFO Continue processing',
    '2023-01-02 ERROR Another error occurred'
]

logs.sort(key=lambda log: log.split()[0]) 
grouped = itertools.groupby(logs, key=lambda log: log.split()[0])

for date, group in grouped:
    print(f'Date: {date}')
    for log in group:
        print(f'    {log}')

输出：

Date: 2023-01-01
    2023-01-01 INFO Start processing
    2023-01-01 ERROR Something went wrong
Date: 2023-01-02
    2023-01-02 INFO Continue processing
    2023-01-02 ERROR Another error occurred