python - Python中的平均计算

Question

我正在尝试加快 python 代码片段的速度。

给定两个大小相等的 (numpy) 数组，目标是找到一个数组(例如 a)中与另一个数组(例如 b)的值相对应的值的平均值。数组的索引是同步的。

例如；

a = np.array([1, 1, 1, 2, 2, 2])
b = np.array([10, 10, 10, 20, 20, 20])

a 中有两个不同的值，1 和 2。b 中的值，其中 a 中的“1”位于相同的索引是 [10, 10, 10]。因此，average(1) 为 10。类似地，average(2) 为 20。

我们可以假设a中的不同值集是先验已知的。 a 中的值不必是连续的，顺序是随机的。我选择这样的例子只是为了便于描述。

这是我的处理方法:

   # Accumulate the total sum and count
    for index, val in np.ndenumerate(a):
        val_to_sum[val] += b[index]
        val_to_count[val] += 1

    # Calculate the mean
    for val in val_to_sum.keys():
        if val_to_count[val]:  # skip vals with zero count
            val_to_mean[val] = val_to_sum[val] / val_to_count[val]

这里 val_to_sum 和 val_to_count 是字典，它们根据 a 中可以看到的已知值列表初始化为零(1在本例中为 2)。

我怀疑这是最快的计算方法。我预计列表会很长，比如说几百万，并且可能值的集合约为数十个。

如何加快计算速度？

解决方案可以吗？ 受到以下答案之一的启发，这可能会做到:

for val in a 
  b[a==val].mean() 

Answer 1

也许这样的事情会起作用:

import numpy as np

a = np.array([1, 1, 1, 2, 2, 2])
b = np.array([10, 10, 10, 20, 20, 20])

np.average(b[a==1])
np.average(b[a==2])

对于更大的数据集:

import numpy as np

a = np.random.randint(1,30,1000000)
b = np.random.random(size=1000000)

for x in set(a):
  print("Average for values marked {0}: {1}".format(x,np.average(b[a==x])))