根据条件舍入数组值的快速方法

Question

我有一个像这样的数组：

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

和“条件”数组：

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

如果c[k] - const < a[i][j] < c[k] + const ，我想四舍五入a[i][j]=c[k] ，否则a[i][j] = 0

例如，如果const = 0.05 。 结果可能是：

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

明智的方法是使用3 for循环来检查a[i][j]和c[k]每a[i][j] 。 但是，当a大时，它会非常慢。 我们是否有快速的“ python方式”来做到这一点？

对于循环（慢速）解决方案：

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] 

Answer 1

方法1

一种矢量化方法是broadcasting -

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

样品运行-

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]])

方法＃2

另一个将更接近于我们在问题中所拥有的，但是被向量化了，就像这样-

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

方法3

根据this post ，这是另一个考虑内存效率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