CODEWITHSUNDEEP
arraysnumpypython-3.7
Shew
Shew

Reputation: 1596

Get the row indexes where 1 appears for each column

I have a binary numpy array and the array contains one 1 in some columns or complete zeros i.e. row sum of the array is a binary vector like

A = array([[0, 0, 0, 1, 0, 1],
           [0, 0, 0, 0, 0, 0],
           [0, 1, 0, 0, 0, 0],
           [1, 0, 1, 0, 0, 0]])

I want to find the row indexes where 1 appears in each column. If there is no one in each column then no need to return any index.

In the above case, I want the result as 

[3, 2, 3, 0, 0]

numpy where results are not usable at it iterates row wise

Upvotes: 1

Views: 35

Answers (2)

Paul Panzer
Paul Panzer

Reputation: 53029

You can use where on the transpose of A:

np.where(A.T)[1]
# array([3, 2, 3, 0, 0])

Upvotes: 3

michcio1234
michcio1234

Reputation: 1838

Try this:

rows, cols = np.where(A==1)
result = rows[np.argsort(cols)]

This gives array([3, 2, 3, 0, 0]) as a result. It will only work if your assumption of maximum one appearance of 1 in each column is met.

Upvotes: 1

Related Questions