具有不同位置的 NumPy 索引答案

【问题标题】：NumPy indexing with varying position具有不同位置的 NumPy 索引
【发布时间】：2016-05-17 21:29:29
【问题描述】：

我有一个形状为 (A, B, C) 的数组 input_data 和一个形状为 (B,) 的数组 ind。我想遍历 B 轴并取元素 C[B[i]] 和 C[B[i]+1] 的总和。所需输出的形状为 (A, B)。我有以下有效的代码，但由于基于索引的 B 轴循环，我觉得效率低下。有没有更有效的方法？

import numpy as np

input_data = np.random.rand(2, 6, 10)
ind = [ 2, 3, 5, 6, 5, 4 ]

out = np.zeros( ( input_data.shape[0], input_data.shape[1] ) )

for i in range( len(ind) ):
    d = input_data[:, i, ind[i]:ind[i]+2]
    out[:, i] = np.sum(d, axis = 1)

根据 Divakar 的回答编辑：

import timeit
import numpy as np

N = 1000

input_data = np.random.rand(10, N, 5000)
ind = ( 4999 * np.random.rand(N) ).astype(np.int)

def test_1(): # Old loop-based method
    out = np.zeros( ( input_data.shape[0], input_data.shape[1] ) )

    for i in range( len(ind) ):
        d = input_data[:, i, ind[i]:ind[i]+2]
        out[:, i] = np.sum(d, axis = 1)
    return out

def test_2(): 
    extent = 2 # Comes from 2 in "ind[i]:ind[i]+2"

    m,n,r = input_data.shape
    idx = (np.arange(n)*r + ind)[:,None] + np.arange(extent)
    out1 = input_data.reshape(m,-1)[:,idx].reshape(m,n,-1).sum(2)
    return out1

print timeit.timeit(stmt = test_1, number = 1000)
print timeit.timeit(stmt = test_2, number = 1000)

print np.all( test_1() == test_2(), keepdims = True )

>> 7.70429363482
>> 0.392034666757
>> [[ True]]

【问题讨论】：

标签： python arrays performance numpy vectorization

【解决方案1】：

在broadcasting 的帮助下，这是使用linear indexing 的矢量化方法。我们合并输入数组的最后两个轴，计算与最后两个轴对应的线性索引，执行切片并重新整形为 3D 形状。最后，我们沿最后一个轴求和以获得所需的输出。实现看起来像这样 -

extent = 2 # Comes from 2 in "ind[i]:ind[i]+2"

m,n,r = input_data.shape
idx = (np.arange(n)*r + ind)[:,None] + np.arange(extent)
out1 = input_data.reshape(m,-1)[:,idx].reshape(m,n,-1).sum(2)

如果extent 始终是2 如问题中所述 - "... sum of elements C[B[i]] and C[B[i]+1]"，那么你可以简单地做 -

m,n,r = input_data.shape
ind_arr = np.array(ind)
axis1_r = np.arange(n)
out2 = input_data[:,axis1_r,ind_arr] + input_data[:,axis1_r,ind_arr+1]

【讨论】：

非常感谢！我根据您的回答编辑了问题。
@StewartHolmes 可爱！您也可以尝试第二种方法，假设范围为2。

【解决方案2】：

您也可以将integer array indexing 与basic slicing 结合使用：

import numpy as np

m,n,r = 2, 6, 10
input_data = np.arange(2*6*10).reshape(m, n, r)
ind = np.array([ 2, 3, 5, 6, 5, 4 ])
out = np.zeros( ( input_data.shape[0], input_data.shape[1] ) )
for i in range( len(ind) ):
    d = input_data[:, i, ind[i]:ind[i]+2]
    out[:, i] = np.sum(d, axis = 1)


out2 = input_data[:, np.arange(n)[:,None], np.add.outer(ind,range(2))].sum(axis=-1)
print(out2)
# array([[  5,  27,  51,  73,  91, 109],
#        [125, 147, 171, 193, 211, 229]])

assert np.allclose(out, out2)

【讨论】：