Apply numpy nonzero row-wise?

Question

I have a 2d boolean array from which I'm trying to extract the indices of the true values. Numpy's nonzero function decomposes my 2d array into a list of x's and y's of positions, which is problematic.

Is it possible to find the column indices of the true elements while preserving the row order?

Each of the true values in the columns are associated with each other in the same row so splitting them into (row index, column index) pairs isn't helpful. Is this possible?

I was thinking that maybe np.apply_along_axis maybe useful.

Answer 1

I did not quite understand what you wanted (maybe an example would help), but two guesses:

If you want to see if there are any Trues on a row, then:

np.any(a, axis=1)

will give you an array with boolean value for each row.

Or if you want to get the indices for the Trues row-by-row, then

testarray = np.array([
    [True, False, True],
    [True, True, False],
    [False, False, False],
    [False, True, False]])

collists = [ np.nonzero(t)[0] for t in testarray ]

This gives:

>>> collists
[array([0, 2]), array([0, 1]), array([], dtype=int64), array([1])]

If you want to know the indices of columns with a True on row 3, then:

>>> collists[3]
array([1])  

There is no pure array-based way of accomplishing this because the number of items on each row varies. That is why we need the lists. On the other hand, the performance is decent, I tried it with a 10000 x 10000 random boolean array, and it took 774 ms to complete the task.