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

Question

假设我们有

一个n维numpy.array一个
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 转换为一个元组来获得所需的结果，但是这肯定不是 proper/idiomatic 的方法吗？

>>> 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 的 post.

基于 map 的方法

案例 #1：暗淡 = 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：暗淡 = 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：暗淡 = 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

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

但是不完全numpy基础。

性能问题。 根据以下评论更新。

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

import numpy as np
import timeit

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

imps = "from __main__ import np,a,b"
reps = 100000

tup_cas_t = timeit.Timer("a[tuple(b)]", imps).timeit(reps)
map_rav_t = timeit.Timer("a[map(np.ravel, b)]", imps).timeit(reps)
fla_rp1_t = timeit.Timer("np.take(a,np.ravel_multi_index(b, a.shape))", imps).timeit(reps)
fla_rp2_t = timeit.Timer("a.flat[np.ravel_multi_index(b, a.shape)]", imps).timeit(reps)

print tup_cas_t/map_rav_t  ## 0.505382211881
print tup_cas_t/fla_rp1_t  ## 1.18185817386
print tup_cas_t/fla_rp2_t  ## 1.71288705886

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])

然而，正如您可以从前面的 link 中了解的那样，直接使用数组 b 索引数组 a 意味着您仅使用整个 b 数组索引 a 的第一个索引，这可能会导致输出混乱。

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

Indexing numpy array by a numpy array of coordinates

python

arrays

idioms

numpy

基准测试