Pandas DataFrame按首字母分组筛选多列间隔n行符合条件的行

以下是简化的DataFrame示例：

import pandas as pd
d = {'col1': ['a1', 'a2', 'a3', 'b1', 'b2', 'b3', 'c1', 'c2', 'c3', 'd1', 'd2', 'd3'], 
     'col2': [1, 1, 1, -1, -1, -1, -1, 1, 1, 1, 1, 1], 
     'col3': [-1, -1, 1, -1, -1, 1, 1, 1, 1, -1, 1, 1]}
df = pd.DataFrame(d)

df输出内容：

col1  col2  col3
0    a1     1    -1
1    a2     1    -1
2    a3     1     1
3    b1    -1    -1
4    b2    -1    -1
5    b3    -1     1
6    c1    -1     1
7    c2     1     1
8    c3     1     1
9    d1    -1    -1
10   d2     1    -1
11   d3     1     1

需求说明 #

• 按col1的首字母对行分组
• 筛选规则：每组中col2首次等于1的位置之后，间隔n行首次出现col3等于1的对应行

示例效果：

当n=1（间隔1行）时，输出结果为：

col1  col2  col3
0   d3     1     1

原因：d组col2首次为1的行是d2（索引10），间隔1行的位置是索引11（d3），刚好是该组col3首次为1的行，其余分组不满足条件。

当n=2（间隔2行）时，输出结果为：

col1  col2  col3
0   a3     1     1

原因：a组col2首次为1的行是a1（索引0），间隔2行的位置是索引2（a3），刚好是该组col3首次为1的行，其余分组不满足条件。

核心思路是按首字母分组后，分别定位每组的两个关键位置，直接判断索引差即可，逻辑清晰易维护：

def get_target_rows(df, n):
    result = []
    # 按col1首字母分组
    for _, group in df.groupby(df['col1'].str[0]):
        # 定位当前组col2首次为1的索引
        first_col2_1 = group[group['col2'] == 1].index.min()
        if pd.isna(first_col2_1):
            continue # 组内无col2=1的行，直接跳过
        # 定位当前组col2首次为1之后，col3首次为1的索引
        first_col3_1 = group[(group.index >= first_col2_1) & (group['col3'] == 1)].index.min()
        if pd.isna(first_col3_1):
            continue # 无符合条件的col3=1行，直接跳过
        # 索引差等于n即符合间隔n行的要求
        if first_col3_1 - first_col2_1 == n:
            result.append(group.loc[first_col3_1])
    return pd.DataFrame(result).reset_index(drop=True)

测试调用：

# 测试n=1
print(get_target_rows(df, 1))
# 测试n=2
print(get_target_rows(df, 2))

该方案避免了多层shift嵌套的复杂逻辑，边界情况兼容完善，可读性和可维护性远优于原实现。

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