通过 ndarray 索引 ndarray

Question

我有第一个 ndarray，foo，我想在其中 select 几个元素。

foo = array([0, 10, 30] , [20, 40, 60], [30, 50, 70])

准确地说，我有另一个 ndarray，bar，我在第一个 ndarray 的每一列中存储我想要的行。

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

我想要的结果是：

array([20, 50, 30] , [0, 10, 60])

它是矢量化的方式吗？

当我尝试 foo[bar] 时，它增加了数组的大小。那不是我要找的。

Answer 1

您还需要分别指定每个索引对应的列。

试试这个：

import numpy as np

foo = np.array([[0, 10, 30], [20, 40, 60], [30, 50, 70]])
bar = np.array([[1, 2, 0], [0, 0, 1]])

foo[bar, range(len(foo))]

输出：

array([[20, 50, 30],
       [ 0, 10, 60]])

Answer 2

In [17]: foo[bar, np.arange(3)]
Out[17]: 
array([[20, 50, 30],
       [ 0, 10, 60]])

一维数组np.arange(3)广播到和bar一样的形状所以它相当于

In [35]: X, Y = np.broadcast_arrays(bar, np.arange(3)); Y
Out[35]: 
array([[0, 1, 2],
       [0, 1, 2]])

X 与 bar 相同，因为广播不会改变 bar 的形状。

然后NumPy integer array indexing rules说foo[X, Y]的(i,j)元素等于

foo[X, Y][i, j] = foo[X[i,j], Y[i,j]]

例如，

foo[bar, np.arange(3)][0, 1] = foo[ bar[0,1],  Y[0,1] ]
                             = foo[2, 1]
                             = 50

Indexing ndarray by ndarray