python - 使用 2 维数组切片 3 维数组

Question

假设我们有两个矩阵:

x = np.random.randint(10, size=(2, 3, 3))
idx = np.random.randint(3, size=(2, 3))

问题是使用idx访问x的元素，方式如下:

dim1 = x[0, range(0,3), idx[0]]  # slicing x[0] using idx[0]
dim2 = x[1, range(0,3), idx[1]]
res = np.vstack((dim1, dim2))

有什么巧妙的方法可以做到这一点吗？

Answer 1

您可以将其索引为 basic way ，只是索引器数组的大小必须匹配。这就是那些 .reshape 的用途:

x[np.array([0,1]).reshape(idx.shape[0], -1), 
  np.array([0,1,2]).reshape(-1,idx.shape[1]),
  idx]
Out[29]: 
array([[ 0.10786251,  0.2527514 ,  0.11305823],
       [ 0.67264076,  0.80958292,  0.07703623]])