Extract specific elements from numpy array by column

Question

I would like to extract specific elements from a 2d-array by index. The index specifies the element in the column.

Example:

14,  7, 30  
44, 76, 65  
42, 87, 11

indices = (0, 1, 2) or (0, 1, 1)

=> result =  [14,76,11] or [14, 76, 65]

I don't want to use any loop, just numpy functions and slicing and stuff. I thought about masking but again I don't know how to generate a mask-2d-array from the indices-array without a direct loop.

score 1 · Accepted Answer · answered Oct 26 '14 at 12:07

You can use row and column index vectors directly:

import numpy as np

A = np.array([[14,  7, 30],
              [44, 76, 65],
              [42, 87, 11]])

print A[[0, 1, 2], range(len(A))]
print A[[0, 1, 1], range(len(A))]

(Since you want exactly one item per column, the column index vector is range(len(A)).)

Output:

[14 76 11]
[14 76 65]

Extract specific elements from numpy array by column

1 Answers1