python - 根据字典值计算 numpy 数组中的向量和

Question

我有一个类似下面的数组，但更大:

array = np.random.randint(6, size=(5, 4))

array([[4, 3, 0, 2],
   [1, 4, 3, 1],
   [0, 3, 5, 2],
   [1, 0, 5, 3],
   [0, 5, 4, 4]])

我还有一本字典，它给出了这个数组中每个值的向量表示:

dict_ = {2:np.array([3.4, 2.6, -1.2]), 0:np.array([0, 0, 0]), 1:np.array([3.9, 2.6, -1.2]), 3:np.array([3.8, 6.6, -1.9]), 4:np.array([5.4, 2.6, -1.2]),5:np.array([6.4, 2.6, -1.2])}

我想计算数组中每一行的向量表示的平均值，但是当值为0时，计算平均值时忽略它(字典显示为0向量)。

例如，对于第一行，它应该平均 [5.4, 2.6, -1.2]、[3.8, 6.6, -1.9] 和 [3.4, 2.6, -1.2]，并给出 [4.2, 3.93, - 1.43]作为输出的第一行。

我想要一个保持相同行结构且具有 3 列的输出(字典中的每个向量都有 3 个值)。

如何以高效的方式做到这一点？我的实际字典有超过 100000 个条目，数组是 100000 x 5000。

Answer 1

为了提高效率，我会将字典转换为数组，然后使用高级索引进行查找:

>>> import numpy as np
>>> 
# create problem
>>> v = np.random.random((100_000, 3))
>>> dict_ = dict(enumerate(v))
>>> arr = np.random.randint(0, 100_000, (100_000, 100))
>>> 
# solve
>>> from operator import itemgetter
>>> lookup = np.array(itemgetter(*range(100_000))(dict_))
>>> lookup[0] = np.nan
>>> result = np.nanmean(lookup[arr], axis=1)

或者应用于OP的例子:

>>> arr = np.array([[4, 3, 0, 2],
...    [1, 4, 3, 1],
...    [0, 3, 5, 2],
...    [1, 0, 5, 3],
...    [0, 5, 4, 4]])
>>> dict_ = {2:np.array([3.4, 2.6, -1.2]), 0:np.array([0, 0, 0]), 1:np.array([3.9, 2.6, -1.2]), 3:np.array([3.8, 6.6, -1.9]), 4:np.array([5.4, 2.6, -1.2]),5:np.array([6.4, 2.6, -1.2])}
>>> 
>>> lookup = np.array(itemgetter(*range(6))(dict_))
>>> lookup[0] = np.nan
>>> result = np.nanmean(lookup[arr], axis=1)
>>> result
array([[ 4.2       ,  3.93333333, -1.43333333],
       [ 4.25      ,  3.6       , -1.375     ],
       [ 4.53333333,  3.93333333, -1.43333333],
       [ 4.7       ,  3.93333333, -1.43333333],
       [ 5.73333333,  2.6       , -1.2       ]])

针对@jpp 方法的计时:

pp:    0.8046 seconds
jpp:  10.3449 seconds
results equal: True

产生时间的代码:

import numpy as np

# create problem
v = np.random.random((100_000, 3))
dict_ = dict(enumerate(v))
arr = np.random.randint(0, 100_000, (100_000, 100))

# solve
from operator import itemgetter
def f_pp(arr, dict_):
    lookup = np.array(itemgetter(*range(100_000))(dict_))
    lookup[0] = np.nan
    return np.nanmean(lookup[arr], axis=1)

def f_jpp(arr, dict_):
    def averager(x):
        lst = [dict_[i] for i in x if i]
        return np.mean(lst, axis=0) if lst else np.array([0, 0, 0])

    return np.apply_along_axis(averager, -1, arr)


from time import perf_counter
t = perf_counter()
r_pp = f_pp(arr, dict_)
s = perf_counter()
print(f'pp:  {s-t:8.4f} seconds')
t = perf_counter()
r_jpp = f_jpp(arr, dict_)
s = perf_counter()
print(f'jpp: {s-t:8.4f} seconds')
print('results equal:', np.allclose(r_pp, r_jpp))