0

I have an m x n numpy array arr, and for each column of arr, I have a given range of rows that I want to access.
I have an n x 1 array vec that describes when this range starts.
The range has some constant duration d.

How can I extract this d x n array of interest efficiently?
Can this be done by clever slicing?

My initial thought was to try something like:

arr = np.tile(np.arange(10),(4,1)).T
vec = np.array([3,4,5,4])
d = 3
vec_2 = vec+d
out = arr[vec:vec2,np.arange(n)]

But this gives the following error:

TypeError: only integer scalar arrays can be converted to a scalar index

The desired output would be the following array:

array([[3, 4, 5, 4],
       [4, 5, 6, 5],
       [5, 6, 7, 6],
       [6, 7, 8, 7])

I could loop over d, but performance is important for this piece of code so I would prefer to vectorize it.

SherylHohman
  • 16,580
  • 17
  • 88
  • 94
algol
  • 399
  • 3
  • 15
  • Please provide a [minimal reproducible example](https://stackoverflow.com/help/minimal-reproducible-example). – Dennis Jan 25 '20 at 23:37
  • Use your starting indices in place of `idx` in the accepted answer from the linked Q&A. Use the transposed version of `arr` as the input. – Divakar Jan 26 '20 at 07:08

1 Answers1

1
In [489]: arr=np.arange(24).reshape(6,4)                                                         
In [490]: vec=np.array([0,2,1,3])

Taking advantage of the recent expansion of linspace to generate several arrays:

In [493]: x = np.linspace(vec,vec+2,3).astype(int)                                               
In [494]: x                                                                                      
Out[494]: 
array([[0, 2, 1, 3],
       [1, 3, 2, 4],
       [2, 4, 3, 5]])
In [495]: arr[x, np.arange(4)]                                                                   
Out[495]: 
array([[ 0,  9,  6, 15],
       [ 4, 13, 10, 19],
       [ 8, 17, 14, 23]])

the column iteration approach:

In [498]: np.stack([arr[i:j,k] for k,(i,j) in enumerate(zip(vec,vec+3))],1)                      
Out[498]: 
array([[ 0,  9,  6, 15],
       [ 4, 13, 10, 19],
       [ 8, 17, 14, 23]])
hpaulj
  • 221,503
  • 14
  • 230
  • 353