# Developers Planet

Stewart Holmes February 2016

### NumPy indexing with varying position

I have an array input_data of shape (A, B, C), and an array ind of shape (B,). I want to loop through the B axis and take the sum of elements C[B[i]] and C[B[i]+1]. The desired output is of shape (A, B). I have the following code which works, but I feel is inefficient due to index-based looping through the B axis. Is there a more efficient method?

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

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

Divakar February 2016

Here's a vectorized approach using `linear indexing` with some help from `broadcasting`. We merge the last two axes of the input array, calculate the linear indices corresponding to the last two axes, perform slicing and reshape back to a 3D shape. Finally, we do summation along the last axis to get the desired output. The implementation would look something like this -

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

If the `extent` is always going to be `2` as stated in the question - `"... sum of elements C[B[i]] and C[B[i]+1]"`, then you could simply do -

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

unutbu February 2016

You could also use integer array indexing combined with 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)

print(out2)
# array([[  5,  27,  51,  73,  91, 109],
#        [125, 147, 171, 193, 211, 229]])

assert np.allclose(out, out2)
``````