Extract nonzero elements of each row of an array

Question

I'm using Python 3.9/NumPy 1.22.

Suppose I have a 3x3 matrix:

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

All elements are integers >= 0. Every row has exactly 2 non-zero integers.

I would like to extract the nonzero integers to produce the 3x2 matrix y such that

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

I feel close when using numpy.apply_along_axis, numpy.squeeze, and/or numpy.where, but I'm missing something.

Answer 1

Since you know that you have the same number of zeros per row, you can safely drop them and reshape:

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

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

Output:

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

---

For fun, if you didn't have the same number of zeros, you could move them to the end, then slice to keep the minimum number of them:

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

m = x!=0
out = np.take_along_axis(x, np.argsort(~m, axis=1),
                         axis=1)[:, :m.sum(axis=1).max()]
# array([[10, 40],
#        [40, 90],
#        [90,  0]])