本文介绍在 numpy 中批量提取二维数组每行非零元素的两种核心方法：一种适用于每行非零元素数量固定（如题设中每行恰好 2 个），可直接布尔索引 + reshape 高效实现；另一种适用于各行非零数不等的通用场景，借助排序移位与动态截取完成对齐。
本文介绍在 numpy 中批量提取二维数组每行非零元素的两种核心方法：一种适用于每行非零元素数量固定（如题设中每行恰好 2 个），可直接布尔索引 + 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]])
✅ 推荐方案：布尔索引 + 自动 reshape（适用于每行非零数恒定）
利用 NumPy 的广播布尔索引能力，x != 0 生成同形布尔掩码，x[x != 0] 拉平后返回所有非零元素（按 C 顺序遍历）。由于已知每行有且仅有 k 个非零值（本例中 k = 2），可直接重塑为 (len(x), k)：
y = x[x != 0].reshape(x.shape[0], -1)
? 原理说明：x[x != 0] 返回一维数组 [10, 40, 40, 90, 10, 90]，长度为 3 × 2 = 6；reshape(3, -1) 自动推导第二维为 2，完美复原目标结构。此方法时间复杂度 O(N)，无显式循环，性能最优。
⚠️ 通用方案：动态对齐（适用于各行非零数不等）
若各行非零元素数量不同（如 [[10,40,0], [0,40,90], [0,0,90]]），直接 reshape 会报错。此时需先将零元素“右对齐”，再统一截取前 max_nonzeros 列：
m = x != 0 # 布尔掩码：True 表示非零 sort_idx = np.argsort(~m, axis=1) # 每行内：False(非零)排前，True(零)排后 x_aligned = np.take_along_axis(x, sort_idx, axis=1) max_nonzeros = m.sum(axis=1).max() # 所有行中最多非零数 y = x_aligned[:, :max_nonzeros]
该方案输出为：
array([[10, 40],
[40, 90],
[90, 0]]) # 注意：末行补零以对齐 —— 若需严格去零，应改用列表推导或 masked array
? 提示：若需严格保持每行仅含非零值（不补零），则无法构成规则二维数组，建议使用 Python 列表推导：[row[row != 0] for row in x]，返回 list[np.ndarray]。
总结
- 固定非零数 → 优先用 x[x!=0].reshape(n_rows, -1)：简洁、向量化、高性能；
- 变长非零数 → 用 np.take_along_axis + argsort(~mask) 对齐：灵活但稍重，注意补零语义；
- 避免 apply_along_axis 或 np.where 循环组合——它们破坏向量化优势，显著降低性能。
始终根据数据先验选择策略：规则结构选 reshape，不规则结构优先考虑是否真需矩阵形式，否则列表/结构化数组更自然。
