Efficiently return the index of the first value satisfying condition in array

In this blog post, we will explore various methods to find the index of the first value in a 1D NumPy array or Pandas numeric series that satisfies a given condition. We’ll examine each approach’s performance and discuss optimizations using Numba.

Attempt 1: Using np.where

One common method is to use np.where, which applies a vectorized function to the entire array. However, this can be slow when the condition is met near the start of the array.

# func(arr) returns a Boolean array
idx = next(iter(np.where(func(arr))[0]), -1)

This approach is often too slow due to the overhead of applying func to the entire array. For example:

np.random.seed(0)
arr = np.random.rand(10**7)
m = 0.9

# Start of array benchmark
%timeit next(iter(np.where(arr > m)[0]), -1)                       # 43.5 ms

Attempt 2: Using np.argmax

Another approach is to use np.argmax, which finds the index of the maximum value in the array. However, this fails when a condition is never met.

arr = np.random.rand(10**7)

assert next(iter(np.where(arr > 0.999999)[0]), -1) == np.argmax(arr > 0.999999)

As seen in the benchmarking section, np.argmax is slightly faster than np.where but still slower than Numba’s approach.

Attempt 3: Using a generator expression

A third approach involves using a generator expression with an explicit loop:

# func(arr) returns a Boolean scalar
idx = next((idx for idx, val in enumerate(arr) if func(arr)), -1)

This method is too slow when the condition is met near the end of the array.

Optimizing with Numba

With Numba, it’s possible to optimize both scenarios. Syntactically, you need only construct a function with a simple for loop:

from numba import njit

@njit
def get_first_index_nb(A, k):
    for i in range(len(A)):
        if A[i] > k:
            return i
    return -1

idx = get_first_index_nb(A, 0.9)

Numba improves performance by JIT compiling code and leveraging CPU-level optimizations.

Generalizing the solution

Since Numba permits functions as arguments, we can generalize this approach to find the nth index where a condition is met for an arbitrary func. Here’s how:

@njit
def get_nth_index_count(A, func, count):
    c = 0
    for i in range(len(A)):
        if func(A[i]):
            c += 1
            if c == count:
                return i
    return -1

@njit
def func(val):
    return val > 0.9

# get index of 3rd value where func evaluates to True
idx = get_nth_index_count(arr, func, 3)

For the nth last value, we can feed the reverse array (arr[::-1]) and negate the result from len(arr) - 1.

Performance benchmarking

# Python 3.6.5, NumPy 1.14.3, Numba 0.38.0

np.random.seed(0)
arr = np.random.rand(10**7)
m = 0.9
n = 0.999999

@njit
def get_first_index_nb(A, k):
    for i in range(len(A)):
        if A[i] > k:
            return i
    return -1

def get_first_index_np(A, k):
    for i in range(len(A)):
        if A[i] > k:
            return i
    return -1

%timeit get_first_index_nb(arr, m)                                 # 375 ns
%timeit get_first_index_np(arr, m)                                 # 2.71 µs
%timeit next(iter(np.where(arr > m)[0]), -1)                       # 43.5 ms
%timeit next((idx for idx, val in enumerate(arr) if val > m), -1)  # 2.5 µs

%timeit get_first_index_nb(arr, n)                                 # 204 µs
%timeit get_first_index_np(arr, n)                                 # 44.8 ms
%timeit next(iter(np.where(arr > n)[0]), -1)                       # 21.4 ms
%timeit next((idx for idx, val in enumerate(arr) if val > n), -1)  # 39.2 ms

As seen above, Numba’s approach is significantly faster than the other methods.

In conclusion, using Numba can efficiently find the index of the first value satisfying a condition in a NumPy array or Pandas numeric series. By leveraging JIT compilation and CPU-level optimizations, we can achieve significant performance improvements over traditional vectorized approaches.

Last modified on 2023-06-28