["python","numpy"]
NumPy selecting specific column index per row by using a list of indexes

I'm struggling to select the specific columns per row of a NumPy matrix. Suppose I have the following matrix which I would call X: [1, 2, 3] [4, 5, 6] [7, 8, 9] I also have a list of column indexes per every row which I would call Y: [1, 0, 2] I need to get the values: [2] [4] [9] Instead of a list with indexes Y, I can also produce a matrix with the same shape as X where every column is a bool / int in the range 0-1 value, indicating whether this is the required column. [0, 1, 0] [1, 0, 0] [0, 0, 1] I know this can be done with iterating over the array and selecting the column values I need. However, this will be executed frequently on big arrays of data and that's why it has to run as fast as it can. I was thus wondering if there is a better solution?
If you've got a boolean array you can do direct selection based on that like so: >>> a = np.array([True, True, True, False, False]) >>> b = np.array([1,2,3,4,5]) >>> b[a] array([1, 2, 3]) To go along with your initial example you could do the following: >>> a = np.array([[1,2,3], [4,5,6], [7,8,9]]) >>> b = np.array([[False,True,False],[True,False,False],[False,False,True]]) >>> a[b] array([2, 4, 9]) You can also add in an arange and do direct selection on that, though depending on how you're generating your boolean array and what your code looks like YMMV. >>> a = np.array([[1,2,3], [4,5,6], [7,8,9]]) >>> a[np.arange(len(a)), [1,0,2]] array([2, 4, 9])