如何在Pandas中按id分组生成满足特定条件的col_3列？

我来帮你搞定这个需求！首先咱们先明确规则：按id分组后，对每个行的col_1值I，只要同组内存在另一个col_1值J满足 I < J <= I+10 且J对应的col_2是1，就把col_3设为1，否则设为0。

先看一下你的示例输入DataFrame：

import pandas as pd

# 构造示例输入
df = pd.DataFrame({
    'id': [1,1,1,1,1,1,1,2,2],
    'col_1': [1.2,7.2,12.1,15.2,16.3,21.1,22.2,3,5],
    'col_2': [0,0,1,0,1,0,0,0,1]
})

方法一：直观遍历法（适合小数据集） #

先定义一个处理单个分组的函数，然后按id分组应用这个函数：

def process_group(group):
    # 提取当前组中所有col_2为1的col_1值（候选J）
    valid_j = group[group['col_2'] == 1]['col_1'].values
    # 遍历每个I，检查是否有符合条件的J
    group['col_3'] = group['col_1'].apply(
        lambda i: 1 if any((j > i) and (j <= i + 10) for j in valid_j) else 0
    )
    return group

# 分组处理并重置索引
result = df.groupby('id').apply(process_group).reset_index(drop=True)

运行后得到的结果和你给出的示例完全一致：

id  col_1  col_2  col_3
0   1    1.2      0      0
1   1    7.2      0      1
2   1   12.1      1      1
3   1   15.2      0      1
4   1   16.3      1      0
5   1   21.1      0      0
6   1   22.2      0      0
7   2    3.0      0      1
8   2    5.0      1      0

方法二：Numpy广播优化法（适合大数据集） #

如果你的数据集比较大，循环遍历会比较慢，用Numpy的广播机制可以大幅提升效率：

def process_group_fast(group):
    col1_vals = group['col_1'].values
    # 筛选出col_2为1的J值
    j_candidates = col1_vals[group['col_2'] == 1]
    # 广播实现批量比较：每个I和所有J做区间判断
    # 得到一个形状为(len(col1_vals), len(j_candidates))的布尔矩阵
    mask = (j_candidates > col1_vals[:, None]) & (j_candidates <= col1_vals[:, None] + 10)
    # 每行只要有一个True，就说明存在符合条件的J，转成int类型
    group['col_3'] = mask.any(axis=1).astype(int)
    return group

result_fast = df.groupby('id').apply(process_group_fast).reset_index(drop=True)

这个方法利用Numpy的向量化操作，避免了Python层面的循环，处理大量数据时速度会快很多，结果和方法一完全相同。

关键逻辑解释 #

两种方法的核心都是：

1. 按id分组，保证只在同组内比较；
2. 提前筛选出同组内所有col_2为1的col_1值作为候选J；
3. 对每个I，检查候选J中是否有落在(I, I+10]区间内的值，有则col_3=1，否则0。

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