本文介绍如何高效提取二维 numpy 数组中每行的非零元素，并在已知每行非零元个数一致的前提下，将其重塑为紧凑的二维结果；同时提供通用方案应对各行非零元数量不等的场景。
本文介绍如何高效提取二维 numpy 数组中每行的非零元素，并在已知每行非零元个数一致的前提下，将其重塑为紧凑的二维结果；同时提供通用方案应对各行非零元数量不等的场景。
在科学计算与数据预处理中，常需从稀疏结构（如含大量零值的矩阵）中快速提取有效数值。当数组每行恰好包含相同数量的非零元素时，可借助布尔索引与 reshape 实现极简、向量化、高性能的提取。
以如下 3×3 整数矩阵为例：
import numpy as np
x = np.array([[10, 40, 0],
[0, 40, 90],
[10, 0, 90]])
目标是得到形状为 (3, 2) 的新数组 y，其中每行仅保留原行的两个非零值，顺序不变：
y = np.array([[10, 40],
[40, 90],
[10, 90]])
核心解法（推荐，适用于每行非零元数量固定）：
利用布尔索引 x != 0 得到扁平化的非零值一维数组，再按行数 len(x) 重塑为二维：
out = x[x != 0].reshape(len(x), -1)
该语句等价于：
- x != 0 → 生成同形布尔掩码；
- x[x != 0] → 按 C 风格行优先顺序展平并筛选，结果是一维数组 [10, 40, 40, 90, 10, 90]；
- .reshape(len(x), -1) → 自动推导列数（此处为 6 // 3 = 2），得到 (3, 2) 矩阵。
✅ 优势：完全向量化、无循环、内存局部性好、代码简洁。
⚠️ 注意事项：
- 此方法严格依赖每行非零元素数量一致（本例为 2）。若不满足，reshape 将报错 ValueError: cannot reshape array of size N into shape (a,b)。
- 布尔索引 x != 0 对整数安全；若含浮点数，建议改用 ~np.isclose(x, 0) 避免精度问题。
进阶方案（各行非零元数量不等）：
当各行非零数不同（如 [2, 2, 1]），无法直接 reshape，但可通过“零值右移 + 截断”模拟对齐：
x = np.array([[10, 40, 0],
[0, 40, 90],
[0, 0, 90]]) # 第三行仅1个非零元
m = x != 0 # 布尔掩码
sort_idx = np.argsort(~m, axis=1) # 每行内：False(非零)排前，True(零)排后
aligned = np.take_along_axis(x, sort_idx, axis=1)
max_nnz = m.sum(axis=1).max() # 最大非零数（此处为2）
out = aligned[:, :max_nnz] # 截取前 max_nnz 列
输出为：
array([[10, 40],
[40, 90],
[90, 0]]) # 注意：末尾补零（非原始非零值），如需纯非零可额外掩码
? 提示：若需严格保持每行仅含非零值（允许不等宽），则必须使用 Python 列表推导或 np.array([row[row != 0] for row in x], dtype=object)，但会失去数值数组的向量化优势。
综上，对于题目所述“每行恰有 2 个非零整数”的典型场景，x[x != 0].reshape(len(x), -1) 是最优雅、高效且可读性强的标准解法。
