Home > front end >  Select non-consecutive row and column indices from 2d numpy array
Select non-consecutive row and column indices from 2d numpy array

Time:01-27

I have an array a

a = np.arange(5*5).reshape(5,5)

array([[ 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 want to select the last two columns from row one and two, and the first two columns of row three and four.

The result should look like this

array([[3,  4, 10, 11],
       [8,  9, 15, 16]])

How to do that in one go without indexing twice and concatenation?

I tried using take

a.take([[0,1,2,3], [3,4,0,1]])
array([[0, 1, 2, 3],
       [3, 4, 0, 1]])

ix_

a[np.ix_([0,1,2,3], [3,4,0,1])]
array([[ 3,  4,  0,  1],
       [ 8,  9,  5,  6],
       [13, 14, 10, 11],
       [18, 19, 15, 16]])

and r_

a[np.r_[0:2, 2:4], np.r_[3:5, 0:2]]
array([ 3,  9, 10, 16])

and a combination of ix_ and r_

a[np.ix_([0,1,2,3], np.r_[3:4, 0:1])]
array([[ 3,  0],
       [ 8,  5],
       [13, 10],
       [18, 15]])

CodePudding user response:

Using integer advanced indexing, you can do something like this

index_rows = np.array([
    [0, 0, 2, 2],
    [1, 1, 3, 3],
])

index_cols = np.array([
    [-2, -1, 0, 1],
    [-2, -1, 0, 1],
])

a[index_rows, index_cols]

where you just select directly what elements you want.

  • Related