引言
在数据分析和处理中，我们经常需要根据特定分组内的逻辑来生成新的数据列。这些逻辑可能涉及排序、条件判断以及累积计算。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 如下所示：
| id | date | date_difference | number | text |
|---|---|---|---|---|
| 1 | 2019-02-01 | NULL | 1 | A |
| 2 | 2019-02-10 | 9 | 0 | A |
| 3 | 2019-02-25 | 15 | 1 | A |
| 4 | 2019-03-05 | 11 | 0 | A |
| 5 | 2019-03-16 | 10 | 0 | A |
| 6 | 2019-04-05 | 19 | 0 | B |
| 7 | 2019-05-15 | 40 | 0 | B |
我们的目标是根据 text 列进行分组，并在每个组内，依据 date 列的降序，生成一个名为 test 的新列。生成 test 列的规则如下：
- 在每个组内，从日期降序排列的第一个条目开始计算。
- 当 number 列的值为 0 时，步长（step size）初始为 1。
- 当遇到 number 列的值为 1 时，步长增加 1。
- 如果组内 number 列中没有 1，则整个组的步长始终保持为 1。
期望的最终 DataFrame 如下：
| id | date | date_difference | number | text | test |
|---|---|---|---|---|---|
| 1 | 2019-02-01 | NULL | 1 | A | 2 |
| 2 | 2019-02-10 | 9 | 0 | A | 2 |
| 3 | 2019-02-25 | 15 | 1 | A | 1 |
| 4 | 2019-03-05 | 11 | 0 | A | 1 |
| 5 | 2019-03-16 | 10 | 0 | A | 1 |
| 6 | 2019-04-05 | 19 | 0 | B | 1 |
| 7 | 2019-05-15 | 40 | 0 | B | 1 |
解决方案实现
解决此问题的关键在于正确地结合 Pandas 的分组、排序、位移和累积求和操作。
核心思路
- 分组 (Group by): 首先，我们需要根据 text 列对 DataFrame 进行分组，因为 test 列的计算逻辑是针对每个 text 组独立的。
- 组内排序 (Sort within group): 问题的关键在于“从日期降序开始计算”。这意味着在每个组内部，我们需要先按 date 列降序排列数据，然后进行计算。
- 位移 (Shift): 为了实现“当 number == 0 时步长为 1，当找到 1 时步长增加 1”的逻辑，我们可以将 number 列进行位移操作。将 number 列向前位移一位，并用 1 填充位移后产生的第一个缺失值。这样，对于原始 number 为 0 的行，其对应的位移后值通常是其下一个条目的 number 值；而对于原始 number 为 1 的行，其对应的位移后值会影响其之前的行的计算。fill_value=1 确保了序列的起始值（即日期降序的第一个值）为 1，符合“步长初始为 1”的条件。
- 累积求和 (Cumulative Sum): 对位移后的 number 列进行累积求和，即可得到所需的 test 值。这个累积和会根据 number 列中的 1 进行递增。
- 结果对齐 (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' 的数据（按原始索引顺序）：
| id | date | number |
|---|---|---|
| 1 | 2019-02-01 | 1 |
| 2 | 2019-02-10 | 0 |
| 3 | 2019-02-25 | 1 |
| 4 | 2019-03-05 | 0 |
| 5 | 2019-03-16 | 0 |
-
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 .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])
-
.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
.droplevel("text"): apply 方法在返回 Series 时，如果 groupby 包含多个键或 apply 的结果不是单个 Series，可能会产生 MultiIndex。在这里，apply 内部返回的是一个 Series，其索引是原始 DataFrame 的索引，但由于 groupby("text")，其上会有一个 text 层的 MultiIndex。droplevel("text") 移除了这个额外的索引层，使得最终 Series 的索引与原始 df 的索引完全匹配。
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），对于大型数据集应谨慎使用。然而，对于这种复杂的、需要组内排序和
