CODEWITHSUNDEEP
pythonnumpy
recodeFuture
recodeFuture

Reputation: 325

Access array elements by array/list of column indexes

I have 2D array of integers (e.g. A with A.shape (10,5)) and 1D array of column indexes (which generally differ among each other) (e.g. idx with idx.shape (10,)). From i-th row, I would like to acquire the element from array A with column index idx[i]. What would be the best (fastest) solution if the desired output is 1D array of acquired elements or list of those elements?

A = np.arange(50).reshape(10,5)
idx=np.array([2,4,0,0,3,1,3,1,2,4])

Desirable output:

output = [2,9,10,15,23,26,33,36,42,49]

or

output = np.array([2,9,10,15,23,26,33,36,42,49])

Upvotes: 0

Views: 1005

Answers (3)

hpaulj
hpaulj

Reputation: 231385

A[np.arange(A.shape[0]), idx]

np.diag(A[:,idx]) works, but does more work than necessary. A[:,idx] is a (10,10)array (in this example), which is then whittled down to a(10,)`.

For this small array, mine is 2x faster. For a (100,50) it is 16x faster.

Upvotes: 2

Jon Clements
Jon Clements

Reputation: 142146

Using numpy you can take the diagonal of a column index, eg:

np.diag(A[:,idx])
# array([ 2,  9, 10, 15, 23, 26, 33, 36, 42, 49])

Upvotes: 5

NDevox
NDevox

Reputation: 4086

You can use enumerate to access the row number you are on, and then use that to access the index you need on idx like so:

output = []

for row in enumerate(A):
    output.append(row[1][idx[row[0]]])

From this I got output == [2, 9, 10, 15, 23, 26, 33, 36, 42, 49]

Upvotes: 2

Related Questions