Pandas快速计算连续相同组最大值

引言：理解连续分组求最大值的挑战

在数据分析中，我们经常需要对数据集中的特定模式进行聚合计算。一个常见的需求是，当某一列中出现连续相同的数值时，我们希望对这些连续的行进行分组，并计算其另一列的聚合值（例如最大值、最小值、平均值等）。

考虑以下Pandas DataFrame：

import pandas as pd

data = {
    'Fruits': ['Apple', 'Apple', 'Banana', 'Orange', 'Apple', 'Apple'],
    'Price': [20, 30, 50, 170, 55, 90]
}

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

输出：

   Fruits  Price
0   Apple     20
1   Apple     30
2  Banana     50
3  Orange    170
4   Apple     55
5   Apple     90

我们的目标是：对于连续出现的相同水果（例如连续的“Apple”），找出它们对应的“Price”列中的最大值，并将这个最大值填充回这些连续行的“Max”列中。具体来说，对于索引0和1的两个连续“Apple”，我们期望其“Max”值为30；对于索引4和5的两个连续“Apple”，我们期望其“Max”值为90。

初学者可能会尝试使用循环或基于shift()的条件判断来解决此问题。例如，以下尝试虽然使用了shift()，但存在明显缺陷：

# 初始尝试（存在问题）
df_attempt = df.copy()
df_attempt['Max'] = 0.0 # 初始化Max列
for x in range(0, df_attempt.shape[0]):
    condition = (df_attempt['Fruits'] == "Apple") & \
                (df_attempt['Fruits'].shift(-1) == "Apple") | \
                (df_attempt['Fruits'].shift(1) == "Apple")
    # 错误：这里计算的是满足条件的所有行的全局最大值，而不是每个连续组的最大值
    df_attempt.loc[condition, 'Max'] = df_attempt.loc[condition, 'Price'].max()

print(df_attempt)

输出：

   Fruits  Price   Max
0   Apple     20  90.0
1   Apple     30  90.0
2  Banana     50  90.0
3  Orange    170   0.0
4   Apple     55  90.0
5   Apple     90  90.0

这种方法未能达到预期效果，主要原因有两点：

1. 条件判断的误区： (df['Fruits'].shift(1)=="Apple") 这样的条件会使得紧跟在“Apple”后面的行（如“Banana”行）也被错误地纳入条件范围。在上述例子中，索引为2的“Banana”行，其前一行（索引1）是“Apple”，因此该条件为真。
2. 聚合范围错误： 最关键的问题在于 df_attempt.loc[condition, 'Price'].max() 这一步。它计算的是所有满足 condition 的行的 Price 列的全局最大值（在这个例子中是90），而不是针对每个独立的连续“Apple”组的最大值。因此，所有被条件命中的行都被赋上了同一个全局最大值90。

核心解决方案：利用 shift() 和 cumsum() 创建连续分组ID

Pandas提供了更高效和准确的方法来解决这类连续分组问题，核心在于巧妙地结合使用shift()、ne()和cumsum()来创建唯一的连续组ID，然后配合groupby()和transform()进行聚合。

1. 识别连续值块并生成组ID

首先，我们需要一种机制来识别数据框中“Fruits”列连续相同值的起始点。shift()函数可以帮助我们获取前一行（或后一行）的值，而ne()（不等于）则可以用来比较当前行与前一行是否相同。

• df.Fruits.shift()：将“Fruits”列向下移动一行，第一行会变为NaN。
• df.Fruits.ne(df.Fruits.shift())：比较当前“Fruits”值是否与前一个“Fruits”值不相等。如果相等，结果为False；如果不相等（意味着新的连续块开始），结果为True。

import pandas as pd

data = {
    'Fruits': ['Apple', 'Apple', 'Banana', 'Orange', 'Apple', 'Apple'],
    'Price': [20, 30, 50, 170, 55, 90]
}
df = pd.DataFrame(data)

# 步骤一：创建连续分组ID
# df.Fruits.shift()       -> [NaN, 'Apple', 'Apple', 'Banana', 'Orange', 'Apple']
# df.Fruits.ne(df.Fruits.shift()) -> [True, False, True, True, True, False]
# (True表示与前一个不同，即新组开始)
grp = df.Fruits.ne(df.Fruits.shift()).cumsum()
print("连续分组ID (grp):\n", grp)

cumsum()（累积求和）函数对这个布尔序列进行操作：每当遇到True时，累积和加1。这样，每个连续的相同值块都会被赋予一个唯一的整数ID。

输出：

连续分组ID (grp):
 0    1
1    1
2    2
3    3
4    4
5    4
Name: Fruits, dtype: int32

可以看到，索引0和1的“Apple”被分到了组1，索引2的“Banana”被分到了组2，索引3的“Orange”被分到了组3，而索引4和5的“Apple”被分到了组4。这正是我们想要的连续分组。

2. 按组应用聚合操作并广播结果

有了这些组ID，我们就可以使用Pandas的groupby()方法进行分组聚合。为了将聚合结果（例如最大值）填充回原始DataFrame的相应行，我们需要使用transform()方法。

• df.groupby(grp)：根据上一步生成的grp序列进行分组。
• ['Price'].transform('max')：在每个分组内部，对“Price”列计算最大值，然后transform()会将这个最大值广播回该组内的所有原始行，确保返回的Series与原DataFrame的索引和长度一致。

完整代码示例

将上述两个步骤结合起来，得到解决问题的完整且高效的代码：

import pandas as pd

data = {
    'Fruits': ['Apple', 'Apple', 'Banana', 'Orange', 'Apple', 'Apple'],
    'Price': [20, 30, 50, 170, 55, 90]
}

df = pd.DataFrame(data)

# 创建连续分组ID
grp = df.Fruits.ne(df.Fruits.shift()).cumsum()

# 根据分组ID计算每个连续块的Price最大值，并使用transform广播结果
df['Max'] = df.groupby(grp)['Price'].transform('max')

print(df)

结果分析与验证

运行上述代码，将得到以下输出：

   Fruits  Price  Max
0   Apple     20   30
1   Apple     30   30
2  Banana     50   50
3  Orange    170  170
4   Apple     55   90
5   Apple     90   90

观察结果，我们可以看到：

• 索引0和1的“Apple”行，其“Max”值均为30，这是它们所在连续组（20, 30）的最大值。
• 索引2的“Banana”行，其“Max”值为50，这是它所在连续组（50）的最大值。
• 索引3的“Orange”行，其“Max”值为170，这是它所在连续组（170）的最大值。
• 索引4和5的“Apple”行，其“Max”值均为90，这是它们所在连续组（55, 90）的最大值。

这完美地满足了我们的需求，并且解决了初始尝试中出现的问题。

注意事项与最佳实践

• 通用性： 这种shift().ne().cumsum()结合groupby().transform()的方法非常通用。它不仅可以用于查找最大值，还可以应用于任何其他聚合函数，如min()、mean()、sum()、count()等，只需替换transform('max')中的聚合函数即可。
• 效率： 相较于使用Python原生的for循环进行行迭代，Pandas的矢量化操作（shift(), ne(), cumsum(), groupby(), transform()）在处理大型数据集时效率要高得多，因为它是在C语言层面进行优化计算的。
• 处理 NaN 值： shift()在默认情况下会在第一个位置产生NaN。ne()会将NaN与实际值视为不相等，这通常是期望的行为，因为它会确保第一个非NaN值始终被视为新组的开始。
• 多列分组： 如果需要根据多列的连续相同值进行分组，可以将这些列组合成一个元组或字符串，然后应用相同的ne().cumsum()逻辑。例如：grp = (df['ColA'] + df['ColB']).ne((df['ColA'] + df['ColB']).shift()).cumsum()。

总结

掌握shift().ne().cumsum()与groupby().transform()的组合是Pandas数据处理中的一个重要技巧。它提供了一种优雅、高效且准确的方法来识别和处理数据框中的连续相同值块，并对其进行各种聚合计算。这种模式在处理时间序列数据、日志数据或任何需要基于连续模式进行分析的场景中都非常有用。