本文介绍如何高效提取二维 numpy 数组中每行的非零元素，并统一组织为形状规整的新矩阵，特别适用于每行非零元数量固定（如恒为 2）的场景，核心方法是布尔索引结合 reshape。
本文介绍如何高效提取二维 numpy 数组中每行的非零元素，并统一组织为形状规整的新矩阵，特别适用于每行非零元数量固定（如恒为 2）的场景，核心方法是布尔索引结合 reshape。
在科学计算和数据预处理中，常需从稀疏结构的数组中剥离有效值。例如，给定一个 3×3 的非负整数矩阵，每行恰好含 2 个非零元素，目标是提取这些非零值，构造一个 3×2 的紧凑矩阵：
import numpy as np
x = np.array([[10, 40, 0],
[0, 40, 90],
[10, 0, 90]])
最简洁、高效的解法是利用 布尔索引 + 自动重塑：
y = x[x != 0].reshape(x.shape[0], -1)
✅ 原理说明：
- x != 0 生成与 x 同形的布尔数组；
- x[x != 0] 拉平（flattened）提取所有 True 位置的值，返回一维数组 [10, 40, 40, 90, 10, 90]；
- reshape(x.shape[0], -1) 将其重排为 3 行、列数自动推导（即 6 // 3 = 2 列），完美匹配需求。
⚠️ 注意事项：
- 此方法严格依赖每行非零元素数量一致（本例为 2）。若不满足，reshape 会报错 ValueError: cannot reshape array of size X into shape (a,b)；
- 不要使用 np.apply_along_axis 或循环，它们在 NumPy 中效率较低，且难以保证输出形状统一；
- 避免 np.where 单独配合索引赋值——虽可行，但代码冗长且易出错。
? 进阶参考（非等长情况）：
若各行非零元数量不同（如 [2, 2, 1]），可先将零移至行尾，再按最少非零数截断（需额外逻辑）：
m = x != 0 # 将 False（零）排到右侧，True（非零）靠左 sorted_x = np.take_along_axis(x, np.argsort(~m, axis=1), axis=1) max_nnz = m.sum(axis=1).max() # 最大非零数（如 2） y_flexible = sorted_x[:, :max_nnz] # 截取前 max_nnz 列
但该方案会引入填充（如补零），且失去原始顺序语义。因此，在已知每行非零数恒定的前提下，首选 x[x!=0].reshape(rows, -1) —— 简洁、向量化、零开销、一行解决。
