Snjór
Snjór

Reputation: 61

Numpy: Efficient way of finding the index of the element in an array with the smallest value given an index array

Say I have a numpy array a = np.array([1, 5, 3, 2, 4, 6, 7]). Now I have another numpy array b = np.array([-1, -2, 3, 2, -1, -3]). The length of b is smaller than or equal to a. I wanna find the index i of the smallest element in a such that b[i] > 0. So in the example above, the result will be 3 since according to b only indices 2, 3 are valid and a[2] == 3 and a[3] == 2, so index 3 is chosen.

My current solution is

    smallest = np.inf
    index = None
    for i in range(len(b)):
        if b[i] > 0:
            if(a[i] < smallest):
                smallest = a[i]
                index = i

I am not sure if I can use numpy to do it more efficiently. Any advice is appreciated. Thank you.

Upvotes: 2

Views: 58

Answers (3)

PiRK
PiRK

Reputation: 1055

one liner:

idx = np.argwhere(a==a[:len(b)][b>0].min())[0]

Understandable code:

shortened_a = a[:len(b)]
filtered_a = shortened_a[b>0]
smallest = filtered_a.min()
indices_of_smallest = np.argwhere(a==smallest)
first_idx = indices_of_smallest[0]

Upvotes: 0

Paritosh Singh
Paritosh Singh

Reputation: 6246

You can use the intermediate results of indices from b to get the right index later, heres a way.

import numpy as np
a = np.array([1, 5, 3, 2, 4, 6, 7])
b = np.array([-1, -2, 3, 2, -1, -3])

indices_to_check = np.where(b > 0)[0]
result = indices_to_check[np.argmin(a[indices_to_check])]
#Output:
3

Upvotes: 0

Divakar
Divakar

Reputation: 221514

Here's one vectorized way -

In [72]: idx = np.flatnonzero(b>0)

In [73]: idx[a[:len(b)][idx].argmin()]
Out[73]: 3

Upvotes: 1

Related Questions