如何用Python 3.9/NumPy 1.22提取数组每行的非零元素？

针对你给出的3×3矩阵，已知每行恰好包含2个非零整数，以下几种NumPy方法可以快速提取出目标3×2矩阵：

方法一：布尔索引+重塑形状（最简洁） #

直接筛选所有非零元素，再利用reshape调整为目标维度：

import numpy as np
x = np.array([[10, 40, 0],
              [0, 40, 90],
              [10, 0, 90]])

y = x[x != 0].reshape(-1, 2)

执行后y即为：

array([[10, 40],
       [40, 90],
       [10, 90]])

这个方法的前提是每行非零元素数量固定为2，这样扁平化后的非零元素总数是3*2=6，reshape(-1,2)会自动匹配行数。

方法二：按行处理（更直观） #

用np.apply_along_axis按行过滤非零元素，之前失败可能是因为返回值的维度处理问题，试这个写法：

def get_non_zero(row):
    return row[row != 0]

y = np.apply_along_axis(get_non_zero, axis=1, arr=x)

apply_along_axis会将每行传入函数，返回的长度为2的一维数组会被自动拼接成3×2的矩阵。

方法三：通过索引定位提取 #

用np.where获取非零元素的位置，再按行索引提取：

# 获取所有非零元素的行、列索引
row_idx, col_idx = np.where(x != 0)
# 将列索引按行分组（每行2个）
col_idx = col_idx.reshape(-1, 2)
# 按行提取对应列的元素
y = x[np.arange(x.shape[0])[:, None], col_idx]

这种方法更灵活，即使后续非零元素位置变化也能准确提取。

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