使用二维数组的行索引进行Numpy高级索引，不使用广播输出

Question

我有一个ndarray数组array，它有三个维度，并且一些包含两个维度的索引ndarray数组idxs，用于指定array的第一维度的索引。 idxs的第一个维度与array的第二个维度匹配，即idxs.shape[0] == array.shape[1]。

我想得到一个结果ndarray数组result，它具有三个维度和形状(idxs.shape[1], array.shape[1], array.shape[2])，如下所示：

for i0 in range(idxs.shape[1]):
    for i1 in range(array.shape[1]):
        result[i0, i1] = array[idxs[i1, i0], i1]

如何更直接地获得这个？

我考虑使用高级索引，但我不确定它应该是什么样子。

在 Theano 中，以下内容有效：

dim1 = theano.tensor.arange(array.shape[1])
result = array[idxs[dim1], dim1]

- Albert

Answer 1

创建一个二维索引网格，对应行索引：idxs[i1, i0]，并使用N x 1数组进行列索引。当像这样索引到array时，列索引将被broadcasted到行索引的形状。因此，我们将采用基于broadcasted indexing的方法，如下所示 -

# Get 2D grid of row indices corresponding to two nested loops
row_idx = idxs[np.arange(array.shape[1])[:,None],np.arange(idxs.shape[1])]

# Use column indices alongwith row_idx to index into array. 
# The column indices would be broadcasted when put as Nx1 array.
result = array[row_idx,np.arange(array.shape[1])[:,None]].T

请注意，如@ali_m的评论中所提到的那样，np.ix_也可以用于创建row_idx，代码如下：

row_idx = idxs[np.ix_(np.arange(array.shape[1]),np.arange(idxs.shape[1]))]

运行时测试和验证输出

函数定义：

def broadcasted_indexing(array,idxs):
    row_idx = idxs[np.arange(array.shape[1])[:,None],np.arange(idxs.shape[1])]
    return array[row_idx,np.arange(array.shape[1])[:,None]].T

def forloop(array,idxs):
    result = np.zeros((idxs.shape[1],array.shape[1]))
    for i0 in range(idxs.shape[1]):
        for i1 in range(array.shape[1]):
            result[i0, i1] = array[idxs[i1, i0], i1]
    return result

运行时测试并验证输出：

In [149]: # Inputs
     ...: m = 500
     ...: n = 400
     ...: array = np.random.rand(m,n)
     ...: idxs = np.random.randint(0,array.shape[1],(n,m))
     ...: 

In [150]: np.allclose(broadcasted_indexing(array,idxs),forloop(array,idxs))
Out[150]: True

In [151]: %timeit forloop(array,idxs)
10 loops, best of 3: 136 ms per loop

In [152]: %timeit broadcasted_indexing(array,idxs)
100 loops, best of 3: 5.01 ms per loop

- Divakar

---

有一个方便的函数np.ix_专门为此目的而设计。 - ali_m

@ali_m 谢谢！我老是忘记这个，现在已经加上了。那么，这可以替换掉如何创建“row_idx”，但似乎np.ix_不能用于二维数组输入，因此最后一步仍需要“广播索引”，对吗？我甚至不确定该怎么称呼它 :) - Divakar

实际上，row_idx与idxs完全相同，因此你可以直接使用 array[idxs, np.arange(array.shape[1])[:, None]].T(或者array[idxs.T, np.r_[:array.shape[1]]]作为紧凑的方式)。 - ali_m

@ali_m 哦，是的！最初我并没有假设 idxs.shape[0] == array.shape[1]，但有了它，似乎 idxs 就与 row_idx 相同了。那个 np.r_ 是另一个很棒的工具，非常感谢！考虑将所有这些发布为答案吧！ - Divakar

Answer 2

你的for循环执行以下操作：

out[i, j] == array[idxs[j, i], j]

换句话说，idxs中的第j,i个元素给出了array中i,j个元素对应的行索引。相应的列索引就是0到idxs.shape[0] - 1之间的整数序列（在这种情况下与array.shape[1] - 1相同，但一般情况下不需要这样）。

因此，您的for循环可以用单个数组索引操作来替换，如下所示：

def simplified(array, idxs):
    return array[idxs.T, np.arange(idxs.shape[0])]

我们可以针对@Divakar答案中的函数进行正确性和速度测试：

m, n = 500, 400
array = np.random.rand(m, n)
idxs = np.random.randint(n, size=(n, m))

print(np.allclose(forloop(array, idxs), simplified(array, idxs)))
# True

%timeit forloop(array, idxs)
# 10 loops, best of 3: 101 ms per loop

%timeit broadcasted_indexing(array, idxs)
# 100 loops, best of 3: 4.1 ms per loop

%timeit simplified(array, idxs)
# 1000 loops, best of 3: 1.66 ms per loop

- ali_m

---

非常感谢提供简化版本。我不太明白为什么这与 array[idxs.T] 相同。它总是尝试匹配多个索引向量，但当维度的索引是隐式的时，却不这样做？ - Albert

这种情况经常让我困惑。理解一维情况要容易得多。idxs.T被解释为行索引的二维数组，所以如果array是一个(m,)的一维数组，那么array[idxs.T]的形状将是(m, n)（因为你从每一行中多次采样）。在你的情况下，array已经是(m, n)，所以array[idxs.T]的结果是(m, n, n)，因为numpy保留了“现有”的列维度。为了折叠“现有”的列维度，你需要给它另一个一维向量的索引。 - ali_m