Question

我有一个ndarray array和ndim 3，还有一些索引ndarray idxs和ndim 2，它们指定array第一维的索引。 idxs的第一维与array的第二维匹配，即idxs.shape[0] == array.shape[1]。

我希望得到一个结果为ndarray result的ndim 3和形状(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]

Answer 1

您的for循环执行此操作：

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

也就是说，idxs中的j，i ^th元素将行索引赋予array i，j < out中的sup> th 元素。 array中相应的列索引集只是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

Answer 2

创建与行索引相对应的索引的2D网格：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

Numpy使用2D数组行索引进行高级索引，而不广播输出

2 个答案: