I have an np array, for example
import numpy as np
X = np.array([[1,2,3],[4,5,6], [7,8,9]])
and another vector, which represents the columns I want to keep for each row:
y = [1,2,0].
Is there a vectorized way to keep only the relevant X[i,j]s? The desired output would be
X_new = np.array([[2],
[6],
[7]])
I came up with few ways to get the desired output. One way would by iterate over rows of X[i,:] and then hstacking over all the rows. A vectorized but highly unreproducible way is to reshape X into an array([1, 2, 3, 4, 5, 6, 7, 8, 9]), transform y to y_new = (i*(ncol+1)) + y) -> y_new=[0+1, 1*3+2, 2*3+0]=[1,5,6], then keep only the X[y_new] and then reshape again.
Is there a simple way to get X_new without all the pyrotechnics?
