如何使用 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。