CODEWITHSUNDEEP
pythonnumpy
Smorts
Smorts

Reputation: 95

Apply numpy 'where' along one of axes

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.

Upvotes: 3

Views: 2827

Answers (2)

Matt Hall
Matt Hall

Reputation: 8152

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

Upvotes: 1

DYZ
DYZ

Reputation: 57075

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

Upvotes: 2

Related Questions