numpy - Numpy：一次索引 2 个维度

Question

我想在某些索引给出的位置上为数组分配新值。一个例子会更清楚：

import numpy as np

#Dimensions
N = 25
n = 50
d = 100
k = 3
p = 7

A = np.random.uniform(size=(N,n,d,d))

A_new_values = np.random.uniform(size=(N,n,k,p))
indexes_new_values = np.random.choice(np.arange(d),size=k*p,replace=False).reshape((k,p))

print(A[:,:,indexes_new_values].shape) 

我希望最后一行 yield (N,n,k,p)，以便能够将新值分配为：

A[:,:,indexes_new_values] = A_new_values

但这会产生错误。我怎样才能分配A_new_values到适当的地方A？

标签： numpyindexing

Answer 1

如果我理解正确，我认为你可以做你想做的事np.put_along_axis：

import numpy as np

#Dimensions
N = 25
n = 50
d = 100
k = 3
p = 7

np.random.seed(0)
A = np.random.uniform(size=(N, n, d, d))
A_new_values = np.random.uniform(size=(N, n, k, p))
indexes_new_values = np.random.choice(np.arange(d), size=k * p, replace=False).reshape((k, p))
np.put_along_axis(A, indexes_new_values.reshape(1, 1, -1, 1), A_new_values.reshape(N, n, -1, 1), axis=2)
print(np.all(A[10, 20, indexes_new_values[1, 5]] == A_new_values[10, 20, 1, 5]))
# True