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): indices = (R[:,0]==i) *( R[:,1]==j) * (R[:,2]==k) return M[indices]

例如

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]

numpy语法来分配数组元素

1 个答案: