Slice and change numpy 2D array with list of colum

Question

This question already has answers here:

Indexing one array by another in numpy (3 answers)

Closed 4 months ago.

I would like to slice a 2D numpy array using a list of column indices. The difficulty is the column indices are different for each row. For example:

x = np.array([[0, 1, 2]
              [3, 4, 5]
              [6, 7, 8]])

and I have a list of column indices as

indices = [[0, 1], [2, 1], [2, 2]]

which means I would like to get column [0,1]for row 0, column [2, 1] for row 1, and column [2, 2] for row 2. The result should be

result = np.array([0, 1],
                  [5, 4],
                  [8, 8]])

How do I use numpy's slicing without involving a for loop?

EDIT:

I saw some answers mentioning using np.take_along_axis(x, indices, 1) This function creates a new array by coping the value from the original array. This is great for reading value from an array, but cannot be used to increment the value of the original array. Is there any similar matrices that can be used for modifying the value of the original array, e.g. np.take_along_axis(x, indices, 1) += 10? Of course the column indices for a row are unique to avoid ambiguity.

Answer 1

You want numpy.take_along_axis

np.take_along_axis(x, np.array(indices), axis = 1)
Out[]: 
array([[0, 1],
       [5, 4],
       [8, 8]])