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): 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]