如何在Pandas中获取滚动最小值对应的类索引列值？

问题描述 #

给定包含ix（类索引列）和val列的DataFrame，需要计算val列滚动窗口（窗口大小为3）最小值对应的ix值，得到目标列ix_min。

示例数据 #

import pandas as pd
import numpy as np

d = pd.DataFrame.from_dict({'ix':[0,1,2,3,4,5], 'val': [1,-2,5,4,9,6]})

方法1：使用rolling.apply实现（直观易读） #

通过rolling.apply遍历每个滚动窗口，找到窗口内val最小值对应的原索引，再提取对应的ix值：

def get_ix_min(window):
    # 获取窗口内val最小值的原索引，返回对应的ix值
    min_row_idx = window.idxmin()
    return d.loc[min_row_idx, 'ix']

# 应用滚动窗口计算，raw=False确保window是带原索引的Series
d['ix_min'] = d['val'].rolling(3).apply(get_ix_min, raw=False)

print(d)

输出结果：

ix  val  ix_min
0   0    1     NaN
1   1   -2     NaN
2   2    5     1.0
3   3    4     1.0
4   4    9     3.0
5   5    6     3.0

方法2：向量化实现（高效适用于大数据集） #

当数据量较大时，apply的遍历效率较低，可采用向量化滑动窗口操作提升性能：

window_size = 3

# 生成所有滚动窗口的索引数组
windows = np.lib.stride_tricks.sliding_window_view(d.index, window_size)
# 找到每个窗口内val最小值的相对位置
min_pos_in_window = d['val'].values[windows].argmin(axis=1)
# 计算对应的ix值
ix_min_values = d['ix'].values[windows[np.arange(len(windows)), min_pos_in_window]]
# 给前window_size-1行填充NaN
d['ix_min'] = np.concatenate([[np.nan]*(window_size-1), ix_min_values])

print(d)

该方法利用numpy的滑动窗口视图实现向量化计算，避免了逐行遍历，性能远高于apply方法。

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