我有一个数组A:
A = array([[1, 2, 3,4], [5,6,7,0] , [8,9,0,0]])
我希望将每行最后一个非零数替换为0。
A = array([[1, 2, 3,0], [5,6,0,0] , [8,0,0,0]])
如何编写代码来创建任意n*m的numpy数组? 谢谢,S ;-)
方法一
一种基于cumsum
和argmax
的方法 -
A[np.arange(A.shape[0]),(A!=0).cumsum(1).argmax(1)] = 0
样例运行 -
In [59]: A
Out[59]:
array([[2, 0, 3, 4],
[5, 6, 7, 0],
[8, 9, 0, 0]])
In [60]: A[np.arange(A.shape[0]),(A!=0).cumsum(1).argmax(1)] = 0
In [61]: A
Out[61]:
array([[2, 0, 3, 0],
[5, 6, 0, 0],
[8, 0, 0, 0]])
方法2
基于 argmax
的另一种方法,希望更加高效 -
A[np.arange(A.shape[0]),A.shape[1] - 1 - (A[:,::-1]!=0).argmax(1)] = 0
解释
argmax
的一个用途是在数组中沿着轴获取 max
元素的第一个出现的 ID。在第一种方法中,我们沿着行累加和并获取第一个 max ID,它代表最后一个非零元素。这是因为在最后一个非零元素之后,剩余元素上的 cumsum
不会增加总和值。
让我们以更详细的方式重新运行该案例 -
In [105]: A
Out[105]:
array([[2, 0, 3, 4],
[5, 6, 7, 0],
[8, 9, 0, 0]])
In [106]: (A!=0)
Out[106]:
array([[ True, False, True, True],
[ True, True, True, False],
[ True, True, False, False]], dtype=bool)
In [107]: (A!=0).cumsum(1)
Out[107]:
array([[1, 1, 2, 3],
[1, 2, 3, 3],
[1, 2, 2, 2]])
In [108]: (A!=0).cumsum(1).argmax(1)
Out[108]: array([3, 2, 1])
A = np.array([[1, 2, 3,4], [5,6,7,0] , [8,9,0,0]])
row_inds = np.arange(A.shape[0])
col_inds = np.arange(A.shape[1])
nonzero_col_inds = (A > 0) * col_inds
max_nonzero_col_inds = np.max(nonzero_col_inds, 1)
A[row_inds, max_nonzero_col_inds] = 0