屏蔽矩阵某些行的少数非零元素

Question

我有一个带有 1 和 0 的 3*3 矩阵 A = [[1,0,1],[0,1,1],[1,0,0]] 和一个指示行和限制的数组 B = [1,2,1]。我想找到 A 的总和超过 B 中相应值的行，并将 A 的非零元素设置为零以确保总和与 B 匹配。查找超过总和的 A 行很容易，但是屏蔽调整总和的元素是我需要帮助的。如何实现（想将其扩展到更大的矩阵和张量）？

Answer 1

我会这样做：

import numpy as np

A = np.array([[1,0,1],[0,1,1],[1,0,0]])
B = np.array([1,2,1])

# a cumulative sum of each row will tell you how many 
# ones were in that row up to each point.
A_cs = np.cumsum(A, axis = 1)
# theresholding according to the sum vector will let 
# you know where you should start omitting values since 
# at that point the sum of the row exceeds its limit.
A_th = A_cs > B[:, None]
# then you can use the boolean array to create a new 
# array where the appropriate values in the original 
# array are set to zero to reduce the row sum.
A_nw = A * (1 - A_th)

输出：

A_nw = 
[[1 0 0]
 [0 1 1]
 [1 0 0]]

无关说明：
以下说明旨在帮助 OP 提高他们与开发相关的搜索技能。 我可以立即回答一些问题，但这不是其中之一。我告诉你这个是因为我通过一个简单的谷歌搜索“python在每一行中找到第i个非零元素”得到了答案，这让我找到了this post，这反过来又让我很快找到了答案。你不必试图成为一个更好、更独立的代码编写者。但是，如果您愿意，请知道您可以。