Question

In the example below I have a 2D array that has some real results that are shifted and padded. The shifts depend on the row (the padding is used to make the array rectangular as required by numpy). Is it possible to extract the real results without a Python loop?

import numpy as np

# results are 'shifted' where the shift depends on the row
shifts = np.array([0, 8, 4, 2], dtype=int)
max_shift = shifts.max()
n = len(shifts)

t = 10 # length of the real results we care about

a = np.empty((n, t + max_shift), dtype=int)
b = np.empty((n, t), dtype=int)

for i in range(n):
    a[i] = np.concatenate([[0] * shifts[i],               # shift
                           (i+1) * np.arange(1, t+1),     # real data
                           [0] * (max_shift - shifts[i])  # padding
                           ])
print "shifted and padded\n", a

# I'd like to remove this Python loop if possible
for i in range(n):
    b[i] = a[i, shifts[i]:shifts[i] + t]

print "real data\n", b

Answer 1

You can use two array to get the data out:

a[np.arange(4)[:, None], shifts[:, None] + np.arange(10)]

or:

i, j = np.ogrid[:4, :10]
a[i, shifts[:, None]+j]

This is called Advanced indexing in NumPy document.