如何加快涉及 pandas 中前一行的计算？

Question

我正在尝试使用自己创建的列的移位值创建一个新的 Pandas DataFrame 列。

我能够做到这一点的唯一方法是遍历数据，这太慢了，导致我的代码出现瓶颈。

import pandas as pd 

df = pd.DataFrame([1,6,2,8], columns=['a'])
df.at[0, 'b'] = 5

for i in range(1, len(df)):
    df.loc[i, ('b')] = (df.a[i-1] + df.b[i-1]) /2

我试过使用 shift，但没用。它为第 1 行填充值，其余为 NaN。我假设此方法无法即时读取新创建的值。

df.loc[1:, ('b')] = (df.a.shift() + df.b.shift()) /2

更新

通过在迭代中使用 df.at 而不是 df.loc

，我能够显着减少时间

def with_df_loc(df):
    for i in range(1, len(df)):
        df.loc[i, ('b')] = (df.a[i-1] + df.b[i-1]) /2
    return df

def with_df_at(df):
    for i in range(1, len(df)):
        df.at[i, 'b'] = (df.a[i-1] + df.b[i-1]) /2
    return df



%timeit with_df_loc(df)
183 ms ± 75.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit with_df_at(df)
19.4 ms ± 2.74 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

此计时基于 150 行的较大数据集。考虑到 df.rolling(20).mean() 大约需要 3 毫秒，我认为这可能是我能做的最好的了。

感谢您的回答，如果我需要进一步优化，我会研究Asish M对numba的建议。

Answer 1

我很可能会误解你的问题，但如果你想创建一个移位的列，试试这个：

df = pd.DataFrame([1,2,3,4], columns=['a'])
df["b"] = df.a.shift()

Answer 2

你可以试试 shift + cumsum, starting from 5 with fillna:

import pandas as pd


df = pd.DataFrame([1,2,3,4], columns=['a'])
df['b'] = df['a'].shift().fillna(5).cumsum()

print(df)

输出

Answer 3

我们可以使用文档中的 numba to speed up calculations here, see Enhancing performance 部分。

import numba 

@numba.njit
def func(a, b_0=5):
    n = len(a)
    b = np.full(n, b_0, dtype=np.float64)
    for i in range(1, n):
        b[i] = (b[i - 1] + a[i - 1]) / 2
    return b

df['b'] = func(df['a'].to_numpy())
df

   a     b
0  1  5.00
1  6  3.00
2  2  4.50
3  8  3.25

性能比较

Benchmarking code, for reference.

蓝线表示您当前方法的最快版本的性能（使用 .at）。橙色线代表 numba 的性能。

如何加快涉及 pandas 中前一行的计算？

How to speed up calculations involving previous row in pandas?

python

dataframe

pandas

numba