使用 numpy 從每個最里面的維度中選擇一個元素

Question

我有一個三維 numpy 源數組和一個二維 numpy 索引數組。

例如：

src = np.array([[[1,2,3],[4,5,6]],
                [[7,8,9],[10,11,12]]])
idx = np.array([[0,1],
                [1,2]])

我想得到一個二維數組，其中每個元素代表 position 中最內層維度的索引值：

array([[1,5],
       [8,12]])

如何使用 numpy 執行此操作？

Answer 1

你可以試試np.take ，這里是文檔。

但是，您應該在展平所有元素后計算數組的索引。 例如，您應該使用

src = np.array([[[1,2,3],[4,5,6]],
                [[7,8,9],[10,11,12]]])
idx = np.array([[0,4],
                [7,11]])

# Wanted result
res = np.take(src, idx)

其中src被視為[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]

你也可以試試np.take_along_axis ，這里是文檔。

使用此方法需要您的src和idx在同一維度，因此，您應該首先解壓縮src並壓縮res 。

# Unsqueezed the last dim
idx = np.expand_dims(idx, axis=-1)

# Squeeze the last dim
res = np.take_along_axis(src, idx, axis=2).squeeze(-1)

Answer 2

您可以使用np.choose方法進行一點重塑：

np.choose(idx.reshape((1, 2, 2)), src.transpose()).reshape((2, 2))

>>>> array([[ 1,  8],
            [ 5, 12]])

Answer 3

直接索引：

src[np.arange(2)[:, None], np.arange(2), idx]