CODEWITHSUNDEEP
pythonnumpy
phsyron
phsyron

Reputation: 530

numpy - select multiple elements from each row of an array

I need to select multiple different values from each row of a 2D array.

A = np.array([[ 1, 2, 3, 4],
              [ 5, 6, 7, 8],
              [ 9,10,11,12])
A[something]

>>> np.array([[ 1, 2],
              [ 6, 7],
              [11,12]])

I know I can create a boolean array the same shape as A and set each element in a for loop, but I'm hoping come up with a better solution.

Upvotes: 5

Views: 7393

Answers (1)

Md Johirul Islam
Md Johirul Islam

Reputation: 5162

You can try the following:

import numpy as np
A = np.array([[ 1, 2, 3, 4],
              [ 5, 6, 7, 8],
              [ 9,10,11,12]])
i = [[0],[1],[2]]
j = [[0,1], [1,2],[2,3]]
B = A[i,j]
print(B)
#Prints
[[ 1  2]
 [ 6  7]
 [11 12]]

Example run

Upvotes: 9

Related Questions