-1

Say I have the following matrix:

[1, 2, 3, 4]
[4, 5 ,6, 7]
[8, 9, 10,11]

So in this example, the column width we desire is 2. and the array size is 3x4. Is there a non loop way of obtaining a desired particular range of columns for each row? Example I want values from columns [2,3] for row 1, [1,2] for row 2 and [0,1] for row 3.

So desired output is :

[3,4]
[5,6]
[8,9]

It would be desirable for the proposed solution to scale to arbitrarily large arrays and larger columns widths.

A similar question was asked here, however, here the person was just looking for a particular column index not a range of column indexes.

zr0gravity7
  • 2,917
  • 1
  • 12
  • 33
Cousik T
  • 1

1 Answers1

0
In [368]: arr = np.arange(1,13).reshape(3,4)
In [369]: idx = np.array([[2,3],[1,2],[0,1]])
In [370]: idx.shape
Out[370]: (3, 2)

Index rows with (3,1), which will broadcast with the (3,2) idx to give:

In [371]: arr[np.arange(3)[:,None], idx]
Out[371]: 
array([[ 3,  4],
       [ 6,  7],
       [ 9, 10]])
In [372]: arr
Out[372]: 
array([[ 1,  2,  3,  4],
       [ 5,  6,  7,  8],
       [ 9, 10, 11, 12]])
hpaulj
  • 221,503
  • 14
  • 230
  • 353