如何在Python中处理Numpy矩阵：按行提取非零值并补0

问题背景 #

给定一个4×18的Numpy矩阵：

import numpy as np

con = np.array([[   0,    0,    0,    0,  500,    0,    0,    0,    0,    0,    0,
                     0,    0,    0,    0,    0,  500,    0],
                 [   0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,
                     0,    0,    0,    0,    0,    0,    0],
                 [   0, 1000,    0,    0,    0,    0,    0,    0,    0,    0,    0,
                   500,    0,    0,    0,    0,    0,    0],
                 [   0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,
                     0,    0,    0,    0,    0,    0,    0]])

需求为：对矩阵的每一行，保留所有大于0的数值；若该行没有大于0的数值，则用[0]填充。

原代码直接将所有符合条件的元素追加到同一个列表，得到一维结果[500, 500, 1000, 500]，但期望得到二维列表[[500, 500], [0], [1000,500], [0]]。

解决方案 #

方法1：改进循环逻辑 #

通过遍历矩阵的每一行，单独处理每行的非零元素，再将处理结果作为子列表加入最终列表：

newArr5 = []
for row in con:
    # 筛选当前行中大于0的元素并转为列表
    non_zero_elements = row[row > 0].tolist()
    # 空列表则用[0]填充，否则用筛选结果
    newArr5.append(non_zero_elements if non_zero_elements else [0])

print(newArr5)
# 输出：[[500, 500], [0], [1000, 500], [0]]

方法2：列表推导式简化写法 #

利用Python的or逻辑（空列表在布尔判断中为False），可以用一行代码实现：

newArr5 = [row[row > 0].tolist() or [0] for row in con]

原代码问题说明 #

原代码的核心问题是未按行分组处理：直接将所有行的符合条件元素追加到同一个列表，导致最终结果是一维结构。改进后的逻辑会为每一行生成一个子列表，从而得到目标二维结构。

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