Is there a numpy alternative to this for loop problem?-CodePudding

I have 3 arrays of the same length:

import numpy as np

weights = np.array([10, 14, 18, 22, 26, 30, 32, 34, 36, 38, 40])
resistances = np.array([15, 16.5, 18, 19.5, 21, 24, 27, 30, 33, 36, 39])
depths = np.array([0,1,2,3,4,5,6,7,8,9,10])

I want to take each item in weights, then find the nearest match that is >= this item in resistances, and then using the index of this nearest match I want to return the corresponding value from depths i.e. depths[index].

BUT, with the additional condition that if nothing is >= the max value in weights then just return last value in depths. I then want to populate a list with the results.

Is there a better way than the for loop approach below? I would like to avoid the loop.

SWP = []
for w in weights:
    if len(depths[w<=resistances]) == 0:
        swp=depths[-1]
    else:
        swp = np.min(depths[w<=resistances])
    SWP.append(swp)
SWP

CodePudding user response：

You can .clip the indices that np.searchsorted produces with len(resistances)-1:

depths[
    np.searchsorted(resistances, weights).clip(max=len(resistances)-1)
]

So any index larger than the last one - will become the last one.

Alternative idea (but only if your resistances are sorted) - clip the weights with the maximum of resistances:

depths[
    np.searchsorted(resistances, weights.clip(max=resistances.max()))
]

CodePudding user response：

Usually to do what you're talking about you want to create a function that can be mapped over a list.

import numpy as np

weights = np.array([10, 14, 18, 22, 26, 30, 32, 34, 36, 38, 40])
resistances = np.array([15, 16.5, 18, 19.5, 21, 24, 27, 30, 33, 36, 39])
depth = np.array([0,1,2,3,4,5,6,7,8,9,10])

def evaluate_weight(w):
    depths = depth[resistances<=w]
    return np.max(depths) if len(depths) else 0

SWP = list(map(evaluate_weight, weights))