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

Question

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?

Do you only have exactly one 1 per row? If not can you provide an updated example? — mozway
– mozway, Commented Dec 16, 2022 at 13:21
No, sometimes there might be rows that have only zeros, because the predictions where lower than a certain threshold, so they didn't become 1. — joasa
– joasa, Commented Dec 16, 2022 at 13:22
Not sure what to do with those predictions actually, maybe the best idea would be to drop them or change them to NaN values, any ideas? — joasa
– joasa, Commented Dec 16, 2022 at 13:24

mozway · Accepted Answer · 2022-12-16 14:02:53Z

2

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.])

edited Dec 16, 2022 at 14:02

answered Dec 16, 2022 at 13:17

mozway

267k13 gold badges55 silver badges106 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

norok2 Over a year ago

perhaps out = np.full(len(preds), np.nan) instead of empty()/fill()?

Collectives™ on Stack Overflow

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

1 Answer 1

handling rows with no match:

1 Comment

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

handling rows with no match:

1 Comment

Related