如何使用 pandas 滚动应用与简单的自定义 function？

Question

我有一个 function func ，我想将其应用于 pandas dataframe 的连续行。 但是，我得到一个ValueError:当我尝试如下操作时。

import numpy as np
import pandas as pd

def func(a: np.ndarray, b: np.ndarray) -> float:
    """Return the sum of sum of vectors a and b"""
    return np.sum(a) + np.sum(b)

df = pd.DataFrame({"a": [1, 2, 3, 4, 5], "b": [10, 11, 12, 13, 14]})
df.rolling(window=2, axis=1).apply(func)
>>>
ValueError: Length of passed values is 2, index implies 5.

我要做的就是将func滚动应用于连续行（这就是我在上面选择window=2的原因）。 下面的代码片段是对此的手动实现。

func(df.iloc[0, :].values, df.iloc[1, :].values)
>>> 24
func(df.iloc[1, :].values, df.iloc[2, :].values)
>>> 28

等等。

请注意，我为func提供的示例仅用于说明目的-我知道在这种情况下您可以使用简单的df.sum(axis=1) + df.shift(-1).sum(axis=1) 。 我想知道的是在一般情况下如何使用滚动申请自定义功能。

Answer 1

我想这可以通过几行代码和一个中间 dataframe 来完成：

import numpy as np
import pandas as pd

def func(a: np.ndarray) -> float:
    return np.sum(a)

df = pd.DataFrame({"a": [1, 2, 3, 4, 5], "b": [10, 11, 12, 13, 14]})
df_rolled = df.rolling(window=2).apply(func)
df["ab_rolled"] = [func([df_rolled["a"][i], df_rolled["b"][i]])
                   for i in range(0, len(df_rolled["a"]))]

print(df)

输出：

   a   b  ab_rolled
0  1  10        NaN
1  2  11       24.0
2  3  12       28.0
3  4  13       32.0
4  5  14       36.0

不过，这口井可能是一个丑陋的代码。 抱歉，这是我第一次使用 pandas。