python - 如何访问滚动运算符中的多列？

Question

我想在需要同时处理两列的 pandas 中做一些滚动窗口计算。我举个简单的例子把问题表述清楚:

import pandas as pd

df = pd.DataFrame({
    'x': [1, 2, 3, 2, 1, 5, 4, 6, 7, 9],
    'y': [4, 3, 4, 6, 5, 9, 1, 3, 1, 2]
})

windowSize = 4
result = []

for i in range(1, len(df)+1):
    if i < windowSize:
        result.append(None)
    else:
        x = df.x.iloc[i-windowSize:i]
        y = df.y.iloc[i-windowSize:i]
        m = y.mean()
        r = sum(x[y > m]) / sum(x[y <= m])
        result.append(r)

print(result)

pandas有没有不用for循环的方法解决问题？感谢任何帮助

Answer 1

您可以使用 rolling window trick for numpy arrays并将其应用于 DataFrame 下的数组。

import pandas as pd
import numpy as np

def rolling_window(a, window):
    shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
    strides = a.strides + (a.strides[-1],)
    return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)

df = pd.DataFrame({
    'x': [1, 2, 3, 2, 1, 5, 4, 6, 7, 9],
    'y': [4, 3, 4, 6, 5, 9, 1, 3, 1, 2]
})

windowSize = 4    

rw = rolling_window(df.values.T, windowSize)
m = np.mean(rw[1], axis=-1, keepdims=True)
a = np.sum(rw[0] * (rw[1] > m), axis=-1)
b = np.sum(rw[0] * (rw[1] <= m), axis=-1)
result = a / b

结果缺少前导 None 值，但它们应该很容易附加(以 np.nan 的形式或在将结果转换为列表之后)。

这可能不是您正在寻找的，与 pandas 一起工作，但它可以在没有循环的情况下完成工作。