numpy 获取二维数组，其中最后一个维度根据二维数组进行索引答案

【问题标题】：numpy get 2d array where last dimension is indexed according to a 2d arraynumpy 获取二维数组，其中最后一个维度根据二维数组进行索引
【发布时间】：2015-05-19 20:12:15
【问题描述】：

我确实阅读了 numpy indexing，但没有找到我要查找的内容。

我有一个 288*384 的图像，其中每个像素都可以在 [0,15] 中进行标记。它存储在一个 3d (288,384,16) 形的 numpy 数组im。

使用im[:,:,1]，例如，我可以获取所有像素都带有标签 1 的图像。

我有另一个二维数组labelling，(288*384) 形，包含每个像素的标签。

如何使用一些巧妙的切片获得每个像素都有对应像素的图像？

使用循环，那就是：

result = np.zeros((288,384))
for x,y in zip(range(288), range(384)):
    result[x,y] = im[x,y,labelling[x,y]]

但这当然是低效的。

【问题讨论】：

有什么进展吗？它起作用了吗？ ...如果您不知道如何接受答案，请参阅here

标签： python numpy multidimensional-array object-slicing

【解决方案1】：

新结果

简短的结果是

np.choose(labelling,im.transpose(2,0,1))

旧结果

试试这个

im[np.arange(288)[:,None],np.arange(384)[None,:],labelling]

适用于以下情况：

import numpy
import numpy.random
import itertools

a = numpy.random.randint(5,size=(2,3,4))
array([[[4, 4, 0, 0],
        [0, 4, 1, 1],
        [3, 4, 4, 2]],

      [[4, 0, 0, 2],
        [1, 4, 2, 2],
        [4, 2, 4, 4]]])

b = numpy.random.randint(4,size=(2,3))
array([[1, 1, 0],
       [1, 2, 2]])

res = a[np.arange(2)[:,None],np.arange(3)[None,:],b]
array([[4, 4, 3],
       [0, 2, 4]])

# note that zip is not doing what you expect it to do
result = np.zeros((2,3))
for x,y in itertools.product(range(2),range(3)):
    result[x,y] = a[x,y,b[x,y]]

array([[4., 4., 3.],
       [0., 2., 4.]])

请注意，zip 没有达到您的预期

zip(range(2),range(3))
[(0, 0), (1, 1)]

可能你的意思是itertools.product

list(itertools.product(range(2),range(3)))
[(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2)]

使用numpy.ix_ 可以避免看起来很可怕的[:,None] 等。

xx,yy = np.ix_( np.arange(2), np.arange(3) )

res = a[xx,yy,b]

【讨论】：