Python Numpy多维数组将负数及后续元素置零的实现方法

你给出的测试用例是不等长子数组，numpy会自动创建object类型数组，逐行处理即可达到预期效果：

import numpy as np 

a = np.array([[12,45,50,60,30],
             [12,34,5,6,88,-10,30,75],
             [3,45,332,45,-12,-4,-64,12],
             [12,45,3,22,323]], dtype=object)

for row in a:
    # 取第一个负数的索引
    neg_idx = np.argmax(row < 0)
    # 排除整行全为正数的场景
    if row[neg_idx] < 0:
        row[neg_idx:] = 0

print(a)

逻辑说明 #

• row < 0会生成和行长度一致的布尔数组，负数对应位置为True
• np.argmax会返回第一个最大值（也就是第一个True）的索引，即第一个负数的位置
• 加判断row[neg_idx] < 0是为了处理全正行：全正行的row < 0全为False，argmax默认返回0，这时候判断第一个元素是否为负，不是则不修改该行
• 确认存在负数后，直接将该索引到行末的所有元素赋值为0即可

运行后输出结果和你预期完全一致：

[[12, 45, 50, 60, 30]
 [12, 34, 5, 6, 88, 0, 0, 0]
 [3, 45, 332, 45, 0, 0, 0, 0]
 [12, 45, 3, 22, 323]]

等长数组优化方案 #

如果你实际使用的是所有子数组长度相同的标准二维numpy数组，可以用向量化操作避免循环，性能更高：

# 假设所有子数组长度均为8
a = np.array([[12,45,50,60,30,0,0,0],
             [12,34,5,6,88,-10,30,75],
             [3,45,332,45,-12,-4,-64,12],
             [12,45,3,22,323,0,0,0]])
# 生成掩码：出现过负数的位置标记为True
mask = np.cumsum(a < 0, axis=1) >= 1
a[mask] = 0

内容的提问来源于stack exchange，提问作者georgehere