Transform a 3D numpy array to 1D based on column value

Maybe this is a very simple task, but I have a numpy.ndarray with shape (1988,3).

preds = [[1 0 0]
        [0 1 0]
        [0 0 0]
        ...
        [0 1 0]
        [1 0 0]
        [0 0 1]]

I want to create a 1D array with shape=(1988,) that will have values corresponding to the column of my 3D array that has a value of 1.

For example,

 new_preds = [0 1 NaN ... 1 0 2]

How can I do this?

Answer

You can use numpy.nonzero:

preds = [[1, 0, 0],
         [0, 1, 0],
         [0, 0, 1],
         [0, 1, 0],
         [1, 0, 0],
         [0, 0, 1]]

new_preds = np.nonzero(preds)[1]

Output: array([0, 1, 2, 1, 0, 2])

handling rows with no match:

preds = [[1, 0, 0],
         [0, 1, 0],
         [0, 0, 0],
         [0, 1, 0],
         [1, 0, 0],
         [0, 0, 1]]

x, y = np.nonzero(preds)

out = np.full(len(preds), np.nan)

out[x] = y

Output: array([ 0., 1., nan, 1., 0., 2.])