问题
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.
回答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]])
来源:https://stackoverflow.com/questions/59666025/slice-and-change-numpy-2d-array-with-list-of-column-indices-different-for-each-r