Fast way to round array values based on condition

Question

I have an array like this:

a = np.array([
     [0.02, 1.01, 4.01, 3.00, 5.12],
     [2.11, 1.50, 3.98, 0.52, 5.01]])

and a "condition" array:

c = np.array([0, 1, 4, 5])

I want to round a[i][j]=c[k] if c[k] - const < a[i][j] < c[k] + const, otherwise a[i][j] = 0

For example, if const = 0.05. The result could be:

a_result = [[0 1 4 0 0]
            [0 0 4 0 5]]

The navie way is to use 3 for loop to check for each a[i][j] and c[k]. However, it's very slow when a is big. Do we have a fast "python way" to do this?

For loop (slow) solution:

a_result = np.full(a.shape, 0)    
const = 0.05     
mh, mw = a.shape
for i in range(mh-1):
   for j in range(mw-1):
      for k in range(1, len(c)):
                if a[i][j] > (c[k] - const) and a[i][j] < (c[k] + const): 
                    a_result[i][j] = c[k]

Check the `map` operator of Python; I think that NumPy has one that is easier to apply to an array. — Prune, Jul 14 '17 at 21:06
How do you make sense of `a[i][j]=c[k]` when `a.shape = (2, 5)` and `c.size = 4`? — jmd_dk, Jul 14 '17 at 21:18
@jmd_dk We do it by imagining it for every element of `a` against every element of `c`. That brings about separate iterators for them. — Divakar, Jul 14 '17 at 21:21
@Divakar OK. You can speed up your loop code somewhat by using `a[i, j]` instead of `a[i][j]`. For dramatic and still easy results, you might want to check out Numba: https://numba.pydata.org/ — jmd_dk, Jul 14 '17 at 21:28

Divakar · Accepted Answer · 2017-07-14T22:12:24.327

Approach #1

One vectorized approach would be with broadcasting -

c[(np.abs(a - c[:,None,None]) < const).argmax(0)]

Sample run -

In [312]: a
Out[312]: 
array([[ 0.02,  1.01,  4.01,  3.  ,  5.12],
       [ 2.11,  1.5 ,  3.98,  0.52,  5.01]])

In [313]: c
Out[313]: array([0, 1, 4, 5])

In [314]: c[(np.abs(a - c[:,None,None]) < const).argmax(0)]
Out[314]: 
array([[0, 1, 4, 0, 0],
       [0, 0, 4, 0, 5]])

Approach #2

Another one that would be closer to what we had in the question, but vectorized, like so -

mask = ((c[:,None,None] - const)  < a) & (a < (c[:,None,None] + const))
out = c[mask.argmax(0)]

Approach #3

Here's another with memory efficiency in mind, based on this post -

idx = np.searchsorted(c, a, side="left").clip(max=c.size-1)
mask = (idx > 0) &  \
     ( (idx == len(xx)) | (np.fabs(yy - xx[idx-1]) < np.fabs(yy - xx[idx])) )
idx0 = idx-mask
out = xx[idx0]
out[np.abs(c[idx0] - a) >= const] = 0

Fast way to round array values based on condition

1 Answers1