How to find the closest value in a list where x > value for every value in a numpy array without looping

user_8067256 0
Viewed 89

For a scalar, let's say I have the following levels

levels_ = [5, 6, 7, 8, 9, 10]

and the function

import numpy as np

def find_level(x, levels):
    levels = np.asarray(levels)
    if x <= levels.min() or x >= levels.max():
        idx = (np.abs(levels - x)).argmin()
    else:
        idx = np.argmax(levels > x) - 1
    return levels[idx]

The function will return the nearest value where x>value from the ordered values of the levels and within the lower and upper boundaries as well.

find_level(x=0.5, levels=levels_) # will return 5
find_level(x=7.01, levels=levels_) # will return 7
find_level(x=7.50, levels=levels_) # will return 7
find_level(x=7.99, levels=levels_) # will return 7
find_level(x=15, levels=levels_) # will return 10

I would like to pass an array x of any size and obtain an array of the same size with the same computation for every value of the array. For example, with the following input:

x_ = np.array([[0.5, 7.01, 7.50, 7.99, 15],
               [20, 30, 6.26, 5.01, 6.54]])

, the function find_level_multi(x=x_, levels=levels_) would the return the following output:

np.array([[5, 7, 7, 7, 10],
          [10, 10, 6, 5, 6]])

I could loop through every row and column and use the function find_level(x, levels) but it would be extremly slow for a large array so I am looking for a function that could accomplish the same task in less time.

multidimensional-array numpy python
2 Answers

Use np.searchsorted to find the indices where the elements in x should be inserted into levels, maintaining the order. Assumes levels is sorted. This solution also works with a scalar or (multi-dimensional) array argument for x:

import numpy as np

def find_level(x, levels):
    """ this assumes levels is sorted """
    levels = np.asanyarray(levels)
    indices = np.searchsorted(levels, np.floor(x))
    indices = np.clip(indices, 0, len(levels) - 1) # clip out-of-bounds index
    return levels[indices]

Testing it:

levels_ = [5, 6, 7, 8, 9, 10]

find_level(x=0.5, levels=levels_) # 5
find_level(x=7.01, levels=levels_) # 7
find_level(x=7.50, levels=levels_) # 7
find_level(x=7.99, levels=levels_) # 7
find_level(x=15, levels=levels_) # 10

x_ = np.array([[0.5, 7.01, 7.50, 7.99, 15],
           [20, 30, 6.26, 5.01, 6.54]])
find_level(x_, levels_)
# array([[ 5,  7,  7,  7, 10],
#        [10, 10,  6,  5,  6]])
user_11702070

You can use broadcasting to build an implicit 3D array from the 2D and perform an np.argmax reduction on one axis. To vectorize the condition, you can use np.select. Here is the resulting vectorized implementation:

def find_level(x, levels):
    levels = np.asarray(levels)
    minIdx = levels.argmin()
    maxIdx = levels.argmax()
    idx = np.argmax(levels[:,None,None] > x[None, :,:], axis=0) - 1
    idx = np.select([x <= levels[minIdx], x >= levels[maxIdx], True], [minIdx, maxIdx, idx])
    return levels[idx]

Here is an example:

x_ = np.array([[0.5, 7.01, 7.50, 7.99, 15],
               [20, 30, 6.26, 5.01, 6.54]])
levels_ = [5, 6, 7, 8, 9, 10]
result = find_level(x_, levels_)
print(result)

Here is the result on the provided input:

[[ 5  7  7  7 10]
 [10 10  6  5  6]]

Note that iterating over all the levels for each cell is probably not the most efficient method if the number of levels is huge.

user_129395570
Related