Question

假设我有两个尺寸相同的数组：一个索引数组idx和一个值数组val：

import numpy as np

idx = np.array([[10, 10, 13],
                [13, 10, 18],
                [10, 16, 18]])

np.random.seed(42)
val = np.round(np.random.uniform(0, 100, 9).reshape((3, 3)), 1)

print(repr(val))
# array([[30.5,  9.8, 68.4],
#        [44. , 12.2, 49.5],
#        [ 3.4, 90.9, 25.9]])

如何使用索引数组中的唯一选择有效地将函数mean应用到值数组上？

这是我最好的for循环的粗暴尝试，其结果是预期的：

{i: val[idx == i].mean() for i in np.unique(idx)}
# {10: 13.975, 13: 56.2, 16: 90.9, 18: 37.7}

虽然在此示例中看起来不错，但在实际情况下，在具有5M元素的数组上具有6000个唯一索引的情况下，它花费了40秒以上。太长了，我正在寻找一种更有效的方法。

Answer 1

对于大型阵列，即使考虑了数据帧设置和系列到字典的转换成本，您也应该发现Pandas更有效：

import numpy as np, pandas as pd
# Python 3.7.0, Pandas 0.23.4, NumPy 1.15.1

np.random.seed(0)
n = 10**3
idx = np.random.randint(0, 20, (n, n))
val = np.random.random((n, n))

df = pd.DataFrame({'idx': idx.ravel(), 'val': val.ravel()})

%timeit pd.DataFrame({'idx': idx.ravel(), 'val': val.ravel()})  # 7.84 ms
%timeit df.groupby('idx')['val'].mean()                         # 34.6 ms
%timeit df.groupby('idx')['val'].mean().to_dict()               # 35.4 ms
%timeit {i: val[idx == i].mean() for i in np.unique(idx)}       # 108 ms

使用另一个数组中的唯一元素将函数应用于值数组

1 个答案: