如何高效提取 DataFrame 中指定起始索引的连续行段（支持单列与多列）

本文介绍一种高效、向量化的方法，替代低效的循环 pd.concat，从原始 DataFrame 中按子集索引批量提取长度为 n 的连续行段，尤其适用于单列选取场景，避免因误用 .iloc 与 [] 操作符导致的索引对齐错误和 NaN 问题。

本文介绍一种高效、向量化的方法，替代低效的循环 `pd.concat`，从原始 dataframe 中按子集索引批量提取长度为 `n` 的连续行段，尤其适用于单列选取场景，避免因误用 `.iloc` 与 `[]` 操作符导致的索引对齐错误和 nan 问题。

在 Pandas 数据处理中，一个常见需求是：给定原始 DataFrame og_df 和一个包含若干起始位置（如某列值或行索引）的子集 sub_df，需提取每个起始位置之后（含自身）连续 n 行的数据，并合并为新 DataFrame。初学者常采用循环 + pd.concat 实现，但该方式存在两大缺陷：时间复杂度呈二次增长（O(k×n)，k 为子集长度），且在单列选取时易因索引未对齐引发意外 NaN。

例如，当对单列执行 og_df['column'].iloc[index:index+n] 并拼接时，Pandas 会尝试按原始 Series 的索引对齐，而 concat 默认按列名/索引合并，导致列名冲突（如生成额外列 0）和缺失值填充，结果完全偏离预期。

✅ 正确解法是避免循环，改用 NumPy 向量化索引构造 + 单次 iloc 切片：

import pandas as pd
import numpy as np

# 示例数据
og_df = pd.DataFrame({'column': range(20), 'other': list('abcdefghijklmnopqrst')})
sub_df = pd.DataFrame({'column': [1, 2, 10]})
n = 3

# ✅ 高效向量化方案
# 步骤1：获取 sub_df['column'] 在 og_df['column'] 中对应的行索引（整数位置）
idx = np.where(og_df['column'].isin(sub_df['column']))[0]  # → array([1, 2, 10])

# 步骤2：为每个起始索引生成 [i, i+1, ..., i+n-1]，展平为一维数组
idx2 = (idx[:, None] + np.arange(n)).ravel()  # → array([1,2,3,2,3,4,10,11,12])

# 步骤3：过滤越界索引，并用 iloc 一次性提取（支持单列或全表）
valid_idx = idx2[idx2 < len(og_df)]
new_df = og_df[['column']].iloc[valid_idx]  # 若需多列，改为 og_df.iloc[valid_idx]

print(new_df)

输出：

⚠️ 关键注意事项：

og_df[['column']] 返回 DataFrame（保留列结构），而 og_df['column'] 返回 Series；若后续需保持 DataFrame 形态（尤其多列场景），务必使用双括号；
np.where(...)[0] 提取的是 位置索引（integer-location），与 iloc 完全兼容；切勿混用 loc 或标签索引；
越界检查 idx2 < len(og_df) 不可省略，否则 iloc 抛出 IndexError；
此方法时间复杂度为 O(N + k×n)，远优于循环 concat 的 O(k²×n)，大数据集性能提升显著。

总结：面对“按条件定位 + 批量滑动窗口提取”任务，应优先采用 NumPy 构造整数索引数组，配合单次 iloc 完成，既保证结果正确性，又兼顾代码简洁性与运行效率。