CODEWITHSUNDEEP
pythonnumpy
a_guest
a_guest

Reputation: 36249

Select different columns for each row

Suppose I have the following array:

>>> a = np.arange(25).reshape((5, 5))
>>> a
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]])

Now I want to select different columns for each row based on the following index array:

>>> i = np.array([0, 1, 2, 1, 0])

This index array denotes the start column for each row and the selections should be of similar range, e.g. 3. Thus I want to obtain the following result:

>>> ???
array([[ 0,  1,  2],
       [ 6,  7,  8],
       [12, 13, 14],
       [16, 17, 18],
       [20, 21, 22]])

I know that I can select a single column per row via

>>> a[np.arange(a.shape[0]), i]

but how about multiple columns?

Upvotes: 0

Views: 633

Answers (1)

akuiper
akuiper

Reputation: 214957

Use advanced indexing with properly broadcasted 2d array as index.

a[np.arange(a.shape[0])[:,None], i[:,None] + np.arange(3)]
#array([[ 0,  1,  2],
#       [ 6,  7,  8],
#       [12, 13, 14],
#       [16, 17, 18],
#       [20, 21, 22]])

idx_row = np.arange(a.shape[0])[:,None]
idx_col = i[:,None] + np.arange(3)

idx_row
#array([[0],
#       [1],
#       [2],
#       [3],
#       [4]])

idx_col
#array([[0, 1, 2],
#       [1, 2, 3],
#       [2, 3, 4],
#       [1, 2, 3],
#       [0, 1, 2]])

a[idx_row, idx_col]
#array([[ 0,  1,  2],
#       [ 6,  7,  8],
#       [12, 13, 14],
#       [16, 17, 18],
#       [20, 21, 22]])

Upvotes: 6

Related Questions