Slice and change numpy 2D array with list of column indices different for each row [duplicate]

江枫思渺然 提交于 2020-01-11 14:40:54

问题


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

易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!