在pandas中取数据帧的特定子集的平均值时，如何获得更好的性能？

Question

我有一个性能问题。下面的代码花了3个小时遍历50000个项目中的5000个项目。

我有一个数据帧df和一个要在key_list中循环的字典键列表。每个键对应于数据帧的单个索引。在每个索引中，我希望获得索引前几行和索引后几行的列的平均值，然后使用mean_cols之前和之后的列创建一个新字典。

mean_cols = ['A', 'B', 'C']
rows_list = []
key_list = list(some_dict.keys()) # around 50k items

for key in key_list:

    means_after = df[mean_cols].iloc[key:key+5].mean()
    means_before = df[mean_cols].iloc[key-5:key].mean()

    for col in mean_cols:
        row_dict[str(col+'_after')] = round(means_after[col], 2)
        row_dict[str(col+'_before')] = round(means_before[col], 2)

    rows_list.append(row_dict)

我很确定是这两条线，

means_after = df[mean_cols].iloc[key:key+5].mean()
means_before = df[mean_cols].iloc[key-5:key].mean()

然而，我想不出更快的方法来做这件事。有谁有什么想法吗？

Answer 1

熊猫.mean()似乎有慢的名声。

我的一个想法是通过使用pandas内置的.to_numpy()进行转换来使用numpy。但是，如果您想要按列计算平均值，numpy的.mean()需要和axis规范-否则它将计算numpy数组中所有值的平均值。

import pandas as pd
import numpy as np
import random

# from @totalhack
mean_cols = ["A", "B"]

df = pd.DataFrame({
    "A": range(0, 50000),
    "B": range(0, 50000)
})

key_list = random.sample(range(50000), k=50000)
# in case that key_list are rownames (indexes), convert them into
# row_indexes, because numpy array won't have names. E.g. by:
# my_rownames = [x for x in your_df_with_rownames.indexes]
# key_list = [my_rownames.index(k) for k in your_old_keylist]

df_mc = np.array(df[mean_cols])

rows_list = []

for key in keys_list:

    means_after = df_mc[key:key+5].mean(axis=0)
    means_before = df_mc[key-5:key].mean(axis=0)
    row_dict = {}

    for col in mean_cols:
        row_dict[str(col+'_after')] = round(means_after[mean_cols.index(col)], 2)
        row_dict[str(col+'_before')] = round(means_before[mean_cols.index(col)], 2)

    rows_list.append(row_dict)

如果数据框只有数字值，那么尽早将其转换为np.arrays将大大加快计算速度。但是，我猜数据框中也有文本或日期数据。所以我想最早转换成mean_cols数组的时间点是在子集.to_numpy()之后--所以我把numpy放在那里。

或者使用并行化(并行使用更多的cpus )

df_mc = np.array(df[mean_cols])

def mean_after(key, np_array=df_mc):
    return list(np.round(np_array[key: key+5].mean(axis=0), 2))

def mean_before(key, np_array=df_mc):
    return list(np.round(np_array[key-5:key].mean(axis=0), 2))


import multiprocessing as mp

pool = mp.Pool()

afters = pool.map(mean_after, keys_list)
befores = pool.map(mean_before, keys_list)

# for what do you need rows_list with dictionaires for each column value?
# why not accessing like this the afters or befores?

afters[row_idx][mean_cols.index(col)]
befores[row_idx][mean_cols.index(col)]

Answer 2

看看pandas的rolling功能。文档列出了各种窗口策略和选项，以尝试解决您的问题。您将使用DataFrame的rolling方法和适合您所需窗口的选项，然后链接一个mean调用以获得结果滚动平均值DataFrame。

编辑:这是一个简单的例子，让你开始。您需要检查输出，并确保这就是您想要的，并且在窗口大小满之前，它也具有NaN值。

如果我将列增加到每个包含50k项，这在我的笔记本电脑上需要7毫秒。

import pandas as pd

mean_cols = ["A", "B"]

df = pd.DataFrame({
    "A": range(0, 50000),
    "B": range(0, 50000)
})

rm_before = df[mean_cols].rolling(5).mean()
rm_after = rm_before.shift(-4)

for col in mean_cols:
    df[col+"_before"] = rm_before[col]
    df[col+"_after"] = rm_after[col]

print(df)