Numpy equivalent of if/else without loop

Approach #1 One approach -

keep_mask = x==50
out = np.where(x>50,0,1)
out[keep_mask] = 50

Approach #2 Alternatively, for in-situ edit -

replace_mask = x!=50
x[replace_mask] = np.where(x>50,0,1)[replace_mask]
# Or (x<=50).astype(int) in place of np.where(x>50,0,1)

Code-golf? If you actually want to play code-golf/one-liner -

(x<=50)+(x==50)*49

Approach #1

For a bit more generic case involving more if-elif parts, we could make use of np.searchsorted -

out_x = np.where(x<=40,0, np.searchsorted([40,50,60,70,80,90], x)+3)

A one-liner that does everything your loops does:

x[x != 50] = x[x != 50] < 50

EDIT:

For your extended question, you'd want something like:

bins = [40, 50, 60, 70, 80, 90, 100]
out = np.digitize(x, bins, right = 1)
out[out.astype(bool)] += 3

np.where(x < 50, 0, 1)

This should be enough. You don't need to keep a mask value for 50 since 50 is neither less than nor greater than 50. Hope this helps.

Related