NumPy数组如何快速将每行第n个非零元素替换为0

完全可以通过NumPy原生向量化操作实现需求，无需编写Python层面的显式循环，执行效率远高于逐行遍历写法，尤其适配宽型二维数组场景。

实现思路 #

• 沿行方向对二进制数组做累计求和，所有值为1的位置的累计计算结果，恰好对应该1在当前行的非零元素排列序号（从1开始计数）
• 生成布尔掩码，标记出累计值刚好等于待删除序号n的位置
• 直接通过布尔索引将标记位置赋值为0即可

代码实现 #

import numpy as np

# 替换为你实际持有的二维二进制数组即可
# np_var = 你的目标数组

# 配置：要替换为0的非零元素序号，从1开始计数
remove_pos = 3

# 核心向量化操作，无显式循环
row_cumsum = np_var.cumsum(axis=1)
# 纯0/1二进制数组可省略后面的&(np_var==1)判断，加了可兼容非严格二进制的数值场景
del_mask = (row_cumsum == remove_pos) & (np_var == 1)
np_var[del_mask] = 0

扩展场景适配 #

如果需要每行删除的非零元素位置不统一（比如第一行删第2个、第二行删第5个），只需要把待删除的序号整理为长度等于数组行数的一维数组remove_pos_arr，将掩码判断逻辑修改为如下形式即可，同样不需要循环：

del_mask = (row_cumsum == remove_pos_arr.reshape(-1, 1)) & (np_var == 1)

效果验证 #

执行完成后可通过如下代码校验每行非零元素数量，确认结果符合预期：

print(np_var.sum(axis=1))
# 输出每行为7，即每行保留7个非零元素

内容的提问来源于stack exchange，提问作者C. Cooney