基于另一列值确定的区间计算最大值/均值（Pandas实现）

原始数据 #

生成目标DataFrame的代码如下：

import pandas as pd
import numpy as np

np.random.seed(5)
df = pd.DataFrame(np.random.randint(10, size=(20, 1)), columns=['A'])
s = [0,0,1,0,0,0,2,0,1,2,0,0,1,0,0,2,0,1,0,0]
df['B'] = s

需求说明 #

需要新增C_Max和D_Mean两列，规则如下：

• 仅对B=1的行进行计算填充
• 以当前B=1的行为起点，取到**下一个B=2的行（包含两端）**的区间
• 将该区间内A列的最大值填入当前行的C_Max，均值填入D_Mean
• 若B=1之后没有出现B=2，则对应的C_Max和D_Mean保留为NaN
• 其他所有行的这两列均为NaN

解决方案代码 #

# 初始化新列为NaN
df['C_Max'] = np.nan
df['D_Mean'] = np.nan

# 提取所有B=1和B=2的索引列表
ones_indices = df[df['B'] == 1].index.tolist()
twos_indices = df[df['B'] == 2].index.tolist()

# 遍历每个B=1的索引，匹配后续最近的B=2
for idx in ones_indices:
    # 筛选出当前B=1之后的所有B=2索引
    valid_twos = [t for t in twos_indices if t > idx]
    if valid_twos:
        # 取第一个符合条件的B=2索引作为区间终点
        end_idx = valid_twos[0]
        # 提取区间内的A列数据
        interval_data = df.loc[idx:end_idx, 'A']
        # 填充统计值
        df.loc[idx, 'C_Max'] = interval_data.max()
        df.loc[idx, 'D_Mean'] = interval_data.mean()

# 查看结果
print(df)

输出结果 #

运行后得到的结果与期望一致：

A  B  C_Max  D_Mean
0   4  0    NaN     NaN
1   1  0    NaN     NaN
2   6  1    6.0     3.4
3   3  0    NaN     NaN
4   4  0    NaN     NaN
5   3  0    NaN     NaN
6   1  2    NaN     NaN
7   4  0    NaN     NaN
8   2  1    3.0     2.5
9   3  2    NaN     NaN
10  4  0    NaN     NaN
11  9  0    NaN     NaN
12  4  1    6.0     4.0
13  0  0    NaN     NaN
14  6  0    NaN     NaN
15  6  2    NaN     NaN
16  9  0    NaN     NaN
17  2  1    NaN     NaN
18  9  0    NaN     NaN
19  3  0    NaN     NaN

内容的提问来源于stack exchange，提问作者Sun Jar