NumPy indexing with varying position
<p>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?</p>

<pre><code>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)
</code></pre>

<p>Edited based on Divakar's answer:</p>

<pre><code>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 )

&gt;&gt; 7.70429363482
&gt;&gt; 0.392034666757
&gt;&gt; [[ True]]
</code></pre>