Pandas按组筛选目标行上下n行触发TypeError的解决求助

给定如下Pandas DataFrame：

import pandas as pd

data = [['A', False, 2], ['A', True, 8], ['A', False, 25], ['A', False, 30], ['B', False, 4], ['B', False, 8], ['B', True, 2], ['B', False, 3]]
df = pd.DataFrame(data=data, columns=['group', 'indicator', 'val'])

初始DataFrame输出：

group  indicator  val
0     A      False    2
1     A       True    8
2     A      False   25
3     A      False   30
4     B      False    4
5     B      False    8
6     B       True    2
7     B      False    3

需求：针对每个group，筛选出indicator == True的行及其上下n行（示例中n=1，即A组保留索引0、1、2的行，B组保留索引5、6、7的行）。

尝试的代码：

# subset each group to list
dfs = [x for _, x in df.groupby('group')] 

for i in dfs:
    # select dataframe
    df_sub = dfs[1]
    # get index of row with indicator True
    idx = df_sub.index[df_sub['indicator'] == True]
    # select n rows above and below row with True
    df_sub = df_sub.iloc[idx - 1: idx + 1]
    # combine each dataframe again
    df_merged = pd.concat(df_sub)
    
    print(df_merged)

触发错误：

TypeError: cannot do positional indexing on Int64Index with these indexers [Int64Index([5], dtype='int64')] of type Int64Index

期望输出：

data = [['A', False, 2], ['A', True, 8], ['A', False, 25], ['B', False, 8], ['B', True, 2], ['B', False, 3]]
df_desired = pd.DataFrame(data=data, columns=['group', 'indicator', 'val'])

对应的输出：

group  indicator  val
0     A      False    2
1     A       True    8
2     A      False   25
3     B      False    8
4     B       True    2
5     B      False    3

1. 循环逻辑错误：循环中硬编码取dfs[1]（第二个分组），没有遍历每个分组，应该使用当前循环变量i作为当前分组的DataFrame。
2. 索引类型混淆：idx = df_sub.index[df_sub['indicator'] == True]获取的是原始DataFrame的标签索引（Int64Index），但iloc需要的是分组内的相对位置索引（整数位置，比如0、1、2），直接用标签索引做iloc切片会触发类型错误。
3. concat用法错误：pd.concat(df_sub)传入的是单个DataFrame，而concat需要接收可迭代对象（如DataFrame列表）；同时循环中每次都会覆盖df_merged，无法保留多个分组的结果。

以下提供两种可行的解决方法：

方法1：使用groupby+apply处理分组 #

通过自定义函数处理每个分组，找到目标行的相对位置后扩展上下n行，同时处理边界情况（比如分组第一行无法取上一行）：

import pandas as pd

data = [['A', False, 2], ['A', True, 8], ['A', False, 25], ['A', False, 30], ['B', False, 4], ['B', False, 8], ['B', True, 2], ['B', False, 3]]
df = pd.DataFrame(data=data, columns=['group', 'indicator', 'val'])

n = 1

def filter_group(group):
    # 获取分组内indicator为True的行的相对位置
    true_positions = group[group['indicator']].index.map(group.index.get_loc)
    # 收集需要保留的位置
    keep_pos = set()
    for pos in true_positions:
        # 从max(0, pos-n)到pos+n+1（iloc左闭右开）
        for p in range(max(0, pos - n), pos + n + 1):
            keep_pos.add(p)
    # 返回筛选后的分组
    return group.iloc[list(keep_pos)]

# 应用到每个分组并合并结果
final_df = df.groupby('group').apply(filter_group).reset_index(drop=True)
print(final_df)

方法2：基于原始索引直接筛选 #

先找到所有indicator == True的行的索引，然后扩展上下n行，同时保证扩展后的索引属于同一个分组：

import pandas as pd

data = [['A', False, 2], ['A', True, 8], ['A', False, 25], ['A', False, 30], ['B', False, 4], ['B', False, 8], ['B', True, 2], ['B', False, 3]]
df = pd.DataFrame(data=data, columns=['group', 'indicator', 'val'])

n = 1

# 获取所有indicator为True的行的原始索引
true_indices = df[df['indicator']].index
keep_indices = set()

for idx in true_indices:
    current_group = df.loc[idx, 'group']
    # 遍历上下n个偏移量
    for offset in range(-n, n+1):
        check_idx = idx + offset
        # 检查索引是否存在且属于同一分组
        if check_idx in df.index and df.loc[check_idx, 'group'] == current_group:
            keep_indices.add(check_idx)

# 按索引排序并生成结果
final_df = df.loc[sorted(keep_indices)].reset_index(drop=True)
print(final_df)

两种方法都能得到符合预期的输出。

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