根据行条件对矩阵进行子集化的高效 Python 解决方案

Question

我正在使用 python 3.9/numpy 1.22。

假设我有一个 3x3 矩阵：

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

所有元素都是> 0的整数。 每行恰好有 2 个非零整数。

我想生成 3x2 矩阵 y 使得

y = np.array([[10,40],[40,90],[10,90]])

使用 numpy.apply_along_axis()、numpy.squeeze() 和/或 numpy.where() 时我感觉很接近，但我错过了一些东西。

Answer 1

由于您知道每行有相同数量的零，因此您可以安全地删除它们并

reshape

：

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

out = x[x!=0].reshape(len(x), -1)

输出：

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