How in numpy get elements of matrix between two indices arrays?

Question

Let's say I have a matrix:

>> a = np.arange(25).reshape(5, 5)`
>> a
[[ 0  1  2  3  4]
 [ 5  6  7  8  9]
 [10 11 12 13 14]
 [15 16 17 18 19]
 [20 21 22 23 24]]

and two vectors of indices that define a span of matrix elements that I want to extract:

>> indices1 = np.array([0, 1, 1, 0, 0])
>> indices2 = np.array([2, 3, 3, 2, 2])

As you can see, difference between each corresponding index is equal to 2.

I would like to do sth like this extract a part of the matrix:

>> submatrix = a[indices1:indices2, :]

so that the result would be 2x5 matrix:

>> submatrix
[[ 0  6  7  3  4],
 [ 5 11 12  8  9]]

For all I know, numpy allows to provide indices as a boundaries, but does not allow to provide arrays, only integers, e.g. a[0:2].

Note what I want to subtract is not a submatrix:

Do you know of some other way of indexing a numpy matrix so that it is possible to provide arrays defining spans? For now I managed to do it only with for loops.

With constant length slices you could construct an indexing array by broadcasting the start array against a [[0],[1]] array. — hpaulj
Are the differences between indices1 and indices2 always a constant? Like it is 2 here? — Divakar
There are lots of very similar Qs on SO. For example, here is a variant with 2D slices from a few days ago stackoverflow.com/q/56207710/7207392 — Paul Panzer
@Divakar They are, but i guess it would be better if there was solution that do not require such constraint and returns all elements (but of course not as a matrix) — dankal444
@PaulPanzer I did my best, but they are hard to find. I would never find the one you posted, because it uses different words to state question and is missing tags like indexing. — dankal444

hpaulj hpaulj · Accepted Answer · 2019-05-22T16:40:37

For reference, the most obvious loop (still took several experimental steps):

In [87]: np.concatenate([a[i:j,n] for n,(i,j) in enumerate(zip(indices1,indices2))], ).reshape(-1,2).T   
Out[87]: 
array([[ 0,  6,  7,  3,  4],
       [ 5, 11, 12,  8,  9]])

Broadcasted indices taking advantage of the constant length:

In [88]: indices1+np.arange(2)[:,None]                                                                   
Out[88]: 
array([[0, 1, 1, 0, 0],
       [1, 2, 2, 1, 1]])
In [89]: a[indices1+np.arange(2)[:,None],np.arange(5)]                                                   
Out[89]: 
array([[ 0,  6,  7,  3,  4],
       [ 5, 11, 12,  8,  9]])

How in numpy get elements of matrix between two indices arrays?

1 Answers