使用 Pandas 高效处理分组数据：基于条件和日期排序创建新列

本文详细介绍了如何利用 Pandas 库处理复杂的分组数据操作。我们将学习如何结合 groupby、apply、sort_values、shift 和 cumsum 等方法，根据特定条件（如日期降序和数值变化）为 DataFrame 添加新列。教程将通过一个实际案例，演示如何高效地实现基于组内逻辑的条件累积计算，并确保结果正确对齐到原始数据结构。

引言

在数据分析和处理中，我们经常需要根据特定分组内的逻辑来生成新的数据列。这些逻辑可能涉及排序、条件判断以及累积计算。Pandas 提供了强大且灵活的工具集来应对此类挑战。本教程将以一个具体的案例为例，展示如何在一个 DataFrame 中，根据 text 列进行分组，然后根据 date 列的降序以及 number 列的数值变化，计算并添加一个名为 test 的新列。

问题描述

假设我们有以下 Pandas DataFrame：

import pandas as pd import numpy as np  data = {     'id': [1, 2, 3, 4, 5, 6, 7],     'date': ['2019-02-01', '2019-02-10', '2019-02-25', '2019-03-05', '2019-03-16', '2019-04-05', '2019-05-15'],     'date_difference': [None, 9, 15, 11, 10, 19, 40],     'number': [1, 0, 1, 0, 0, 0, 0],     'text': ['A', 'A', 'A', 'A', 'A', 'B', 'B'] }  df = pd.DataFrame(data)  print("原始 DataFrame:") print(df)

原始 DataFrame 如下所示：

我们的目标是根据 text 列进行分组，并在每个组内，依据 date 列的降序，生成一个名为 test 的新列。生成 test 列的规则如下：

• 在每个组内，从日期降序排列的第一个条目开始计算。
• 当 number 列的值为 0 时，步长（step size）初始为 1。
• 当遇到 number 列的值为 1 时，步长增加 1。
• 如果组内 number 列中没有 1，则整个组的步长始终保持为 1。

期望的最终 DataFrame 如下：

解决方案实现

解决此问题的关键在于正确地结合 Pandas 的分组、排序、位移和累积求和操作。

核心思路

1. 分组 (Group by): 首先，我们需要根据 text 列对 DataFrame 进行分组，因为 test 列的计算逻辑是针对每个 text 组独立的。
2. 组内排序 (Sort within group): 问题的关键在于“从日期降序开始计算”。这意味着在每个组内部，我们需要先按 date 列降序排列数据，然后进行计算。
3. 位移 (Shift): 为了实现“当 number == 0 时步长为 1，当找到 1 时步长增加 1”的逻辑，我们可以将 number 列进行位移操作。将 number 列向前位移一位，并用 1 填充位移后产生的第一个缺失值。这样，对于原始 number 为 0 的行，其对应的位移后值通常是其下一个条目的 number 值；而对于原始 number 为 1 的行，其对应的位移后值会影响其之前的行的计算。fill_value=1 确保了序列的起始值（即日期降序的第一个值）为 1，符合“步长初始为 1”的条件。
4. 累积求和 (Cumulative Sum): 对位移后的 number 列进行累积求和，即可得到所需的 test 值。这个累积和会根据 number 列中的 1 进行递增。
5. 结果对齐 (Align Results): groupby().apply() 返回的结果通常会保留原始 DataFrame 的索引，但如果内部进行了排序，则需要确保最终结果能正确地与原始 DataFrame 对齐。assign() 方法在这里非常有用，它能自动根据索引将新生成的 Series 与原 DataFrame 合并。

完整代码

import pandas as pd import numpy as np  data = {     'id': [1, 2, 3, 4, 5, 6, 7],     'date': ['2019-02-01', '2019-02-10', '2019-02-25', '2019-03-05', '2019-03-16', '2019-04-05', '2019-05-15'],     'date_difference': [None, 9, 15, 11, 10, 19, 40],     'number': [1, 0, 1, 0, 0, 0, 0],     'text': ['A', 'A', 'A', 'A', 'A', 'B', 'B'] }  df = pd.DataFrame(data)  # 将 'date' 列转换为 datetime 类型，以便正确排序 df['date'] = pd.to_datetime(df['date'])  # 使用 assign 方法添加新列 'test' df_result = df.assign(     test=df     # 1. 按 'text' 列进行分组     .groupby("text")     # 2. 对每个组应用一个函数     .apply(         lambda g: (             # 3. 在组内按 'date' 列降序排序             g.sort_values(by="date", ascending=False)             # 4. 对 'number' 列进行位移，向前一位，并用 1 填充缺失值             .number.shift(periods=1, fill_value=1)             # 5. 对位移后的结果进行累积求和             .cumsum()         )     )     # 6. 移除 apply 产生的 'text' 索引层，使 Series 索引与原始 df 索引一致     .droplevel("text")     # assign 方法会自动将结果 Series 与原始 DataFrame 的索引对齐 )  print("n最终 DataFrame:") print(df_result)

逻辑解析与示例跟踪

让我们以 text 为 ‘A’ 的组为例，详细解释每一步的操作：

原始 text=’A’ 的数据（按原始索引顺序）：

1. g.sort_values(by=”date”, ascending=False): 将组内数据按 date 降序排列。 排序后的数据（索引为原始 DataFrame 索引）：

   id	date	number
   5	2019-03-16	0
   4	2019-03-05	0
   3	2019-02-25	1
   2	2019-02-10	0
   1	2019-02-01	1

2. .number.shift(periods=1, fill_value=1): 对排序后的 number 列 [0, 0, 1, 0, 1] 进行向前位移，并用 1 填充第一个位置。 位移后的 Series: [1, 0, 0, 1, 0] (对应索引 [5, 4, 3, 2, 1])

3. .cumsum(): 对位移后的 Series [1, 0, 0, 1, 0] 进行累积求和。 累积和结果: [1, 1, 1, 2, 2] (对应索引 [5, 4, 3, 2, 1])

   这意味着：

   • df.loc[5, ‘test’] = 1
   • df.loc[4, ‘test’] = 1
   • df.loc[3, ‘test’] = 1
   • df.loc[2, ‘test’] = 2
   • df.loc[1, ‘test’] = 2

4. .droplevel(“text”): apply 方法在返回 Series 时，如果 groupby 包含多个键或 apply 的结果不是单个 Series，可能会产生 MultiIndex。在这里，apply 内部返回的是一个 Series，其索引是原始 DataFrame 的索引，但由于 groupby(“text”)，其上会有一个 text 层的 MultiIndex。droplevel(“text”) 移除了这个额外的索引层，使得最终 Series 的索引与原始 df 的索引完全匹配。

5. df.assign(test=…): assign 方法将这个计算好的 Series 作为 test 列添加到原始 df 中。Pandas 会自动根据索引进行对齐，确保 test 值回到其原始的行位置。

最终结果与期望输出完全一致。对于 text=’B’ 的组，由于 number 列中没有 1，shift(fill_value=1) 会使所有值变为 1，cumsum() 结果也都是 1，符合“没有 1 时步长保持为 1”的规则。

注意事项与总结

• 日期类型转换: 在进行日期排序之前，确保 date 列的数据类型是 datetime。如果不是，需要使用 pd.to_datetime() 进行转换。
• groupby().apply() 的灵活性: apply() 方法非常强大，允许你在每个分组上执行几乎任何自定义操作。但需要注意其性能可能不如优化的 Pandas 方法（如 transform 或 agg），对于大型数据集应谨慎使用。然而，对于这种复杂的、需要组内排序和