[英]How can I update a numpy array with index in another numpy array
I have我有
what I want to do is like the code below我想做的就像下面的代码
for i in range(n1):
for j in range(n2):
for k in range(i1):
b[i, j, k, :] = a[i, j, idx[i, j, k], :]
if there is a numpy function to achieve this without for loop?如果有一个 numpy function 来实现这个不用for循环?
Using as starting point:使用作为起点:
import numpy as np
n1, n2, n3, n4, i1 = range(2, 7)
a = np.random.randint(10, size=(n1, n2, n3, n4))
idx = np.random.randint(n3, size=(n1, n2, i1))
b = np.zeros_like(a, shape=(n1, n2, i1, n4))
In general you can do the following:一般来说,您可以执行以下操作:
I, J, K = np.ogrid[:n1, :n2, :i1]
b[I, J, K] = a[I, J, idx]
Here the I
J
and K
arrays are the equivalent of the loop variables i
j
and k
.这里的
I
J
和K
arrays 相当于循环变量i
j
和k
。 Their shapes have to be in agreement with the shape of idx
.它们的形状必须与
idx
的形状一致。
In case b
has shape (n1, n2, i1, n4) then you might as well do:如果
b
的形状为 (n1, n2, i1, n4) 那么你不妨这样做:
I, J, _ = np.ogrid[:n1, :n2, :1]
b = a[I, J, idx]
Or alternatively without ogrid
:或者没有
ogrid
:
b = np.take_along_axis(a, idx[...,np.newaxis], axis=2)
Here newaxis
is used to insert a length-1 axis to allow broadcasting.这里
newaxis
用于插入一个长度为 1 的轴以允许广播。 Check out the numpy indexing docs for more info.查看numpy 索引文档以获取更多信息。
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.