python - 用于分配数组元素的 Numpy 语法

Question

我想创建一个数组，其元素是其位置的函数。类似的东西

N = 1000000 
newarray = np.zeros([N,N,N])
for i in range(N):
    for j in range(N):
        for k in range(N):
            newarray[i,j,k] = f(i,j,k)

有没有办法通过删除 for 循环/使用 numpy 语法并行化来提高此操作的速度？

这是 f 函数

def f(i,j,k):
索引 = (R[:,0]==i) *( R[:,1]==j) * (R[:,2]==k)
返回 M[索引]


例如在哪里

R = np.random.randint(0,N,[N,3])
M = np.random.randn(N)*15

并且在实际应用中它们并不是随机的。

Answer 1

您可以使用 at 执行该操作方法np.add :

import numpy as np

np.random.seed(0)
N = 100
R = np.random.randint(0, N, [N, 3])
M = np.random.randn(N) * 15
newarray = np.zeros([N, N, N])
np.add.at(newarray, (R[:, 0], R[:, 1], R[:, 2]), M)

在这种情况下，如果R有任何重复的行，则newarray中的相应值将是M中所有相应值的总和.

编辑:要取平均值而不是重复元素的总和，您可以执行以下操作:

import numpy as np

np.random.seed(0)
N = 100
R = np.random.randint(0, N, [N, 3])
M = np.random.randn(N) * 15
newarray = np.zeros([N, N, N])
np.add.at(newarray, (R[:, 0], R[:, 1], R[:, 2]), M)
newarray_count = np.zeros([N, N, N])
np.add.at(newarray_count, (R[:, 0], R[:, 1], R[:, 2]), 1)
m = newarray_count > 1
newarray[m] /= newarray_count[m]