python - Numpy:通过 z 索引数组索引 3d 数组

Question

假设我有一个数组

a = np.zeros((3, 3, 3))

和一个 z 索引数组

z = np.random.randint(0, 3, (3, 3))

假设 z 是

array([[1, 0, 2],
       [2, 2, 1],
       [1, 1, 0]])

现在我想选择带有坐标的 a 值(从 z 的左上角开始并按行遍历数组。按列也可以.) (0, 0, 1), (0, 1, 0), (0, 2, 2), ... 。粗体值来自 z 数组。

Answer 1

我认为这符合你的要求。

import numpy as np 

a = np.arange(3*3*3).reshape(3,3,3)
z = np.array([[1, 0, 2],
       [2, 2, 1],
       [1, 1, 0]])
i = np.arange(a.shape[0]).repeat(a.shape[0]).reshape(a.shape[0], a.shape[1])
j = i.T
#Should do the same as this. possible more efficient, did not test.
#i, j = np.indices(a[...,0].shape)
print a[i,j,z]