Question

假设我们有

n维numpy.array A
a numpy.array B，dtype = int，形状为（n，m）

如何用B对A进行索引，使得结果是一个形状（m，）的数组，其值取自B列所示的位置？

例如，考虑当B是python列表时执行我想要的代码：

>>> a = np.arange(27).reshape(3,3,3)
>>> a[[0, 1, 2], [0, 0, 0], [1, 1, 2]]
array([ 1, 10, 20])    # the result we're after
>>> bl = [[0, 1, 2], [0, 0, 0], [1, 1, 2]]
>>> a[bl]
array([ 1, 10, 20])   # also works when indexing with a python list
>>> a[bl].shape
(3,)

但是，当B是一个numpy数组时，结果是不同的：

>>> b = np.array(bl)
>>> a[b].shape
(3, 3, 3, 3)

现在，我可以通过将B转换为元组来获得所需的结果，但肯定这不是正确/惯用的方法吗？

>>> a[tuple(b)]
array([ 1, 10, 20])

是否有一个numpy函数来实现相同的功能而不将B转换为元组？

Answer 1

一种替代方案是转换为线性索引，然后使用np.take索引或索引到其展平版本 -

np.take(a,np.ravel_multi_index(b, a.shape))
a.flat[np.ravel_multi_index(b, a.shape)]

自定义np.ravel_multi_index以提升效果

我们可以实现自定义版本来模拟np.ravel_multi_index的行为以提升效果，就像这样 -

def ravel_index(b, shp):
    return np.concatenate((np.asarray(shp[1:])[::-1].cumprod()[::-1],[1])).dot(b)

使用它，可以通过两种方式找到所需的输出 -

np.take(a,ravel_index(b, a.shape))
a.flat[ravel_index(b, a.shape)]

基准

此外，还包含问题中基于tuple的方法和来自@ Kanak帖子的基于map的方法。

案例＃1：dims = 3

In [23]: a = np.random.randint(0,9,([20]*3))

In [24]: b = np.random.randint(0,20,(a.ndim,1000000))

In [25]: %timeit a[tuple(b)]
    ...: %timeit a[map(np.ravel, b)]  
    ...: %timeit np.take(a,np.ravel_multi_index(b, a.shape))
    ...: %timeit a.flat[np.ravel_multi_index(b, a.shape)]
    ...: %timeit np.take(a,ravel_index(b, a.shape))
    ...: %timeit a.flat[ravel_index(b, a.shape)]
100 loops, best of 3: 6.56 ms per loop
100 loops, best of 3: 6.58 ms per loop
100 loops, best of 3: 6.95 ms per loop
100 loops, best of 3: 9.17 ms per loop
100 loops, best of 3: 6.31 ms per loop
100 loops, best of 3: 8.52 ms per loop

案例＃2：dims = 6

In [29]: a = np.random.randint(0,9,([10]*6))

In [30]: b = np.random.randint(0,10,(a.ndim,1000000))

In [31]: %timeit a[tuple(b)]
    ...: %timeit a[map(np.ravel, b)]  
    ...: %timeit np.take(a,np.ravel_multi_index(b, a.shape))
    ...: %timeit a.flat[np.ravel_multi_index(b, a.shape)]
    ...: %timeit np.take(a,ravel_index(b, a.shape))
    ...: %timeit a.flat[ravel_index(b, a.shape)]
10 loops, best of 3: 40.9 ms per loop
10 loops, best of 3: 40 ms per loop
10 loops, best of 3: 20 ms per loop
10 loops, best of 3: 29.9 ms per loop
100 loops, best of 3: 15.7 ms per loop
10 loops, best of 3: 25.8 ms per loop

案例＃3：dims = 10

In [32]: a = np.random.randint(0,9,([4]*10))

In [33]: b = np.random.randint(0,4,(a.ndim,1000000))

In [34]: %timeit a[tuple(b)]
    ...: %timeit a[map(np.ravel, b)]  
    ...: %timeit np.take(a,np.ravel_multi_index(b, a.shape))
    ...: %timeit a.flat[np.ravel_multi_index(b, a.shape)]
    ...: %timeit np.take(a,ravel_index(b, a.shape))
    ...: %timeit a.flat[ravel_index(b, a.shape)]
10 loops, best of 3: 60.7 ms per loop
10 loops, best of 3: 60.1 ms per loop
10 loops, best of 3: 27.8 ms per loop
10 loops, best of 3: 38 ms per loop
100 loops, best of 3: 18.7 ms per loop
10 loops, best of 3: 29.3 ms per loop

因此，在处理更高维度的输入和大数据时，寻找替代方案是有意义的。

Answer 2

另一种符合您需求的替代方案是使用np.ravel

Root.A.AC();
console.log(Root.A.AB);

但不完全基于numpy。

<小时/> 的 性能-担忧。 根据以下评论进行了更新。

尽管如此，你的方法比我的好，但不比@Divakar的任何一个好。

>>> a[map(np.ravel, b)] array([ 1, 10, 20])

Answer 3

您在寻找numpy.ndarray.tolist()吗？

>>> a = np.arange(27).reshape(3,3,3)
>>> bl = [[0, 1, 2], [0, 0, 0], [1, 1, 2]]
>>> b = np.array(bl)
>>> a[b.tolist()]
array([ 1, 10, 20])

或者arrays indexing arrays与列表索引非常相似：

>>> a[np.array([0, 1, 2]), np.array([0, 0, 0]), np.array([1, 1, 2])]
array([ 1, 10, 20])

然而，正如您可以从上一个链接那样，使用数组b索引数组a直接意味着您正在使用整个b数组索引第一个索引，这可能导致输出混乱。

通过numpy坐标数组索引numpy数组

3 个答案:

基准