基于Pandas/Numpy优化DataFrame固定间隔选取A列值的方案咨询

问题背景 #

给定如下DataFrame：

import pandas as pd
import numpy as np

df = pd.DataFrame(data=np.random.rand(100, 3), columns=list('ABC'))

需求为按固定间隔选取A列的值，原实现脚本如下：

n=0
for i in range(0, len(df), 22):
    if i+22+n < len(df):
        print(df.A.iloc[[i+n, i+22+n]])
    else: print(i+n, len(df)-1)
    n+=1

原脚本可正常运行，但希望获得基于Pandas或Numpy的更优、更高效的实现方案。

优化方案 #

一、Pandas 向量化实现 #

核心思路是通过向量化操作批量生成所有需要选取的索引对，避免Python循环中的逐次判断，提升效率：

import pandas as pd
import numpy as np

df = pd.DataFrame(data=np.random.rand(100, 3), columns=list('ABC'))

# 计算总分组数
total_groups = len(df) // 22 + 1
# 生成每组的偏移量（对应原脚本中的n）
offsets = np.arange(total_groups)

# 生成每组的两个索引
idx1 = 22 * np.arange(total_groups) + offsets
idx2 = 22 * (np.arange(total_groups) + 1) + offsets
# 处理最后一组索引超出数据范围的情况
idx2 = np.where(idx2 < len(df), idx2, len(df)-1)

# 遍历输出每组结果
for i1, i2 in zip(idx1, idx2):
    if i1 == i2:
        print(i1, i2)
    else:
        print(df['A'].iloc[[i1, i2]])

二、Numpy 原生数组处理 #

若不需要保留Pandas Series的输出格式，直接操作Numpy数组能获得更高的性能，适合大规模数据集：

import pandas as pd
import numpy as np

df = pd.DataFrame(data=np.random.rand(100, 3), columns=list('ABC'))
a_col = df['A'].values

total_groups = len(a_col) // 22 + 1
offsets = np.arange(total_groups)

idx1 = 22 * np.arange(total_groups) + offsets
idx2 = 22 * (np.arange(total_groups) + 1) + offsets
# 用minimum直接截断超出范围的索引
idx2 = np.minimum(idx2, len(a_col)-1)

# 按格式输出结果
for i1, i2 in zip(idx1, idx2):
    if i1 == i2:
        print(i1, i2)
    else:
        print(f"{i1}     {a_col[i1]}")
        print(f"{i2}    {a_col[i2]}")
        print("Name: A, dtype: float64")
        print()

方案优势 #

• 效率更高：用向量化操作替代Python循环中的逐次判断，数据量越大，性能提升越明显
• 可维护性强：间隔参数（如示例中的22）只需修改一处即可调整整个选取逻辑，便于后续修改
• 逻辑清晰：索引生成逻辑集中处理，避免循环内的复杂条件分支，代码可读性更好

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