Python pandas如何将匹配行及其索引±1的行加入新DataFrame

解决方案 #

更推荐使用pandas向量化操作替代循环append的写法，运行效率会高很多，同时可以直接实现你的需求：

import pandas as pd

hra_cc = HRA['COUNTRYCODE'].values.tolist()
# 第一步：筛选出所有符合条件的行的索引
match_idx = Filtered_Data[Filtered_Data['COUNTRYCODE'].isin(hra_cc)].index
# 第二步：生成所有需要保留的索引：原索引、原索引-1、原索引+1，同时做去重、边界过滤
need_idx = pd.Index(match_idx).union(match_idx - 1).union(match_idx + 1)
need_idx = need_idx.intersection(Filtered_Data.index)
# 第三步：一次性取出所有需要的行
df = Filtered_Data.loc[need_idx].copy()

如果一定要在你原有循环的基础上修改，可以按如下方式调整：

hra_cc = HRA['COUNTRYCODE'].values.tolist()
df=pd.DataFrame()
for dest in hra_cc:
    # 先拿到符合当前国家代码的所有行的索引
    match_rows = Filtered_Data.loc[Filtered_Data['COUNTRYCODE']==dest]
    for idx in match_rows.index:
        # 依次添加当前行、前一行、后一行，提前判断索引是否存在避免报错
        for offset in (-1, 0, 1):
            target_idx = idx + offset
            if target_idx in Filtered_Data.index:
                df = df.append(Filtered_Data.loc[target_idx])
# 去重，避免相邻符合条件的行重复添加同一条数据
df = df.drop_duplicates()

注意说明 #

• 两种写法都会自动过滤超出Filtered_Data索引范围的行，比如匹配到第一行时，不会尝试取索引为-1的无效行
• 去重步骤是为了避免多个匹配行的±1索引重叠，导致同一条行被重复添加
• 第一种向量化写法的运行效率远高于循环写法，数据量越大性能优势越明显

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