Apply numpy 'where' along one of axes

Viewed 517

I have an array like that:

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

I would like to find the last occurance of True for each row of the array. If it was 1d array I would do it in this way:

np.where(array)[0][-1]

How do I do something similar in 2D? Kind of like:

np.where(array, axis = 1)[0][:,-1]

but there is no axis argument in np.where.

2 Answers

Since True is greater than False, find the position of the largest element in each row. Unfortunately, argmax finds the first largest element, not the last one. So, reverse the array sideways, find the first True from the end, and recalculate the indexes:

(array.shape[1] - 1) - array[:, ::-1].argmax(axis=1)
# array([0, 0, 0, 1])

The method fails if there are no True values in a row. You can check if that's the case by dividing by array.max(axis=1). A row with no Trues will have its last True at the infinity :)

array[0, 0] = False
((array.shape[1] - 1) - array[:, ::-1].argmax(axis=1)) / array.max(axis=1)
#array([inf,  0.,  0.,  1.])

I found an older answer but didn't like that it returns 0 for both a True in the first position, and for a row of False.

So here's a way to solve that problem, if it's important to you:

import numpy as np

arr = np.array([[False, False, False],  # -1
                [False, False, True],   #  2
                [True, False, False],   #  0
                [True, False, True],    #  2
                [True, True, False],    #  1
                [True, True, True],     #  2
               ])

# Make an adustment for no Trues at all.
adj = np.sum(arr, axis=1) == 0

# Get the position and adjust.
x = np.argmax(np.cumsum(arr, axis=1), axis=1) - adj

# Compare to expected result:
assert np.all(x == np.array([-1, 2, 0, 2, 1, 2]))

print(x)

Gives [-1 2 0 2 1 2].

Related