深入解析 `DataFrame.groupby` 和 `agg` 的用法及使用场景

- 1. `groupby` 的基本用法
- - 语法：
  - 示例：
- 2. `agg` 的基本用法
- - 语法：
  - 示例：
- 3. `first`、`sum`、`lambda` 的用法
- - 3.1 `first`
  - - 示例：
  - 3.2 `sum`
  - - 示例：
  - 3.3 `lambda`
  - - 示例：
- 4. `reset_index` 的用法
- - 语法：
  - 示例：
- 5. 综合示例：`groupby` + `agg` 的完整用法
- 6. 总结

在数据分析和处理中，pandas 是一个非常强大的工具。其中，DataFrame.groupby 和 agg 是两个常用的功能，它们可以帮助我们对数据进行分组、聚合和转换。本文将详细讲解 groupby 和 agg 的用法，以及如何结合 first、sum、lambda 和 reset_index 来实现复杂的数据处理。

1. `groupby` 的基本用法

groupby 是 pandas 中用于分组数据的函数。它的核心思想是将数据按照某个或某些列的值进行分组，然后对每个分组进行操作。

语法：

DataFrame.groupby(by=None, axis=0, level=None, as_index=True, sort=True, group_keys=True, squeeze=False, observed=False)

by：指定分组的列名或列名列表。
as_index：是否将分组列作为索引，默认为 True。
sort：是否对分组键进行排序，默认为 True。

示例：

假设我们有一个包含学生成绩的 DataFrame：

import pandas as pd

data = {
    'Name': ['Alice', 'Bob', 'Alice', 'Bob', 'Alice'],
    'Subject': ['Math', 'Math', 'Science', 'Science', 'Math'],
    'Score': [85, 90, 88, 92, 95]
}

df = pd.DataFrame(data)
print(df)

输出：

    Name  Subject  Score
0  Alice     Math     85
1    Bob     Math     90
2  Alice  Science     88
3    Bob  Science     92
4  Alice     Math     95

我们可以按照 Name 列进行分组：

grouped = df.groupby('Name')
print(grouped.groups)

输出：

{'Alice': [0, 2, 4], 'Bob': [1, 3]}

2. `agg` 的基本用法

agg 是 groupby 对象的一个方法，用于对分组后的数据进行聚合操作。它可以接受一个函数、字符串（如 'sum'、'mean'）或字典（指定不同列的聚合方式）。

语法：

DataFrame.groupby(...).agg(func, *args, **kwargs)

func：聚合函数，可以是字符串、函数或字典。

示例：

我们继续使用上面的 df，计算每个学生的平均分：

result = df.groupby('Name').agg({'Score': 'mean'})
print(result)

输出：

        Score
Name        
Alice   89.333333
Bob     91.000000

3. `first`、`sum`、`lambda` 的用法

在 agg 中，我们可以使用多种聚合函数来处理数据。以下是常用的几种：

3.1 `first`

first 用于获取每个分组中第一行的值。

示例：

result = df.groupby('Name').agg({'Subject': 'first', 'Score': 'mean'})
print(result)

输出：

        Subject      Score
Name                      
Alice     Math  89.333333
Bob       Math  91.000000

3.2 `sum`

sum 用于计算每个分组中数值列的总和。

示例：

result = df.groupby('Name').agg({'Score': 'sum'})
print(result)

输出：

        Score
Name        
Alice     268
Bob       182

3.3 `lambda`

lambda 是一种匿名函数，可以在 agg 中自定义聚合逻辑。

示例：

假设我们想将每个学生的 Score 转换为字符串，并在前面加上 "Score: "：

result = df.groupby('Name').agg({'Score': lambda x: 'Score: ' + str(x.sum())})
print(result)

输出：

                Score
Name                
Alice  Score: 268.0
Bob    Score: 182.0

4. `reset_index` 的用法

reset_index 用于重置 DataFrame 的索引。当使用 groupby 和 agg 后，分组列会作为索引，如果需要将其还原为普通列，可以使用 reset_index。

语法：

DataFrame.reset_index(level=None, drop=False, inplace=False, col_level=0, col_fill='')

drop：是否丢弃索引列，默认为 False。

示例：

result = df.groupby('Name').agg({'Score': 'mean'}).reset_index()
print(result)

输出：

    Name      Score
0  Alice  89.333333
1    Bob  91.000000

5. 综合示例：`groupby` + `agg` 的完整用法

现在，我们来看一个更复杂的例子，结合 first、sum、lambda 和 reset_index 来处理数据。

假设我们有一个包含边信息的 DataFrame：

data = {
    'edge_key': ['A', 'A', 'B', 'B', 'C'],
    'source': ['X', 'X', 'Y', 'Y', 'Z'],
    'target': ['P', 'P', 'Q', 'Q', 'R'],
    'text_unit_ids': [[1, 2], [3], [4], [5, 6], [7]],
    'weight': [10, 20, 30, 40, 50],
    'description': ['Edge A', 'Edge A', 'Edge B', 'Edge B', 'Edge C'],
    'human_readable_id': ['A1', 'A2', 'B1', 'B2', 'C1'],
    'id': [1, 2, 3, 4, 5]
}

df = pd.DataFrame(data)
print(df)

输出：

  edge_key source target text_unit_ids  weight description human_readable_id  id
0        A      X      P        [1, 2]      10     Edge A              A1   1
1        A      X      P           [3]      20     Edge A              A2   2
2        B      Y      Q           [4]      30     Edge B              B1   3
3        B      Y      Q        [5, 6]      40     Edge B              B2   4
4        C      Z      R           [7]      50     Edge C              C1   5

我们希望对 edge_key 进行分组，并进行以下操作：

source 和 target 取第一行的值。
text_unit_ids 合并为一个列表。
weight 求和。
description 拼接为一个字符串。
human_readable_id 和 id 取第一行的值。

代码如下：

def _aggregate_text_unit_ids(x):
    return [item for sublist in x for item in sublist]

def _add_quotes(s):
    return f'"{s}"'

updated_edges_df = df.groupby('edge_key').agg(
    source=('source', 'first'),
    target=('target', 'first'),
    text_unit_ids=('text_unit_ids', _aggregate_text_unit_ids),
    weight=('weight', 'sum'),
    description=('description', lambda x: ''.join(x.apply(_add_quotes))),
    human_readable_id=('human_readable_id', 'first'),
    id=('id', 'first'),
).reset_index(drop=True)

print(updated_edges_df)

输出：

  source target text_unit_ids  weight description human_readable_id  id
0      X      P     [1, 2, 3]      30  "Edge A""Edge A"              A1   1
1      Y      Q  [4, 5, 6]       70  "Edge B""Edge B"              B1   3
2      Z      R        [7]       50        "Edge C"              C1   5