python - 矩阵向量差的有效元素级argmin

Question

假设一个数组 a.shape == (N, M)和一个向量 v.shape == (N,) .目标是计算 argmin的 abs的 v从 a 的每个元素中减去- 那是，

out = np.zeros(N, M)
for i in range(N):
    for j in range(M):
        out[i, j] = np.argmin(np.abs(a[i, j] - v))

我有一个 vectorized implementation通过 np.matlib.repmat ，而且速度要快得多，但需要 O(M*N^2)内存，在实践中是 Not Acceptable 。计算仍然在 CPU 上完成，所以最好的办法似乎是在 C 中实现 for 循环作为扩展，但也许 Numpy 已经实现了这个逻辑。
可以？任何可以有效实现上述功能的即用型 Numpy 函数？

Answer 1

灵感来自 this post ，我们可以利用 np.searchsorted ——

def find_closest(a, v):
    sidx = v.argsort()
    v_s = v[sidx]
    idx = np.searchsorted(v_s, a)
    idx[idx==len(v)] = len(v)-1
    idx0 = (idx-1).clip(min=0)
    
    m = np.abs(a-v_s[idx]) >= np.abs(v_s[idx0]-a)
    m[idx==0] = 0
    idx[m] -= 1
    out = sidx[idx]
    return out

性能更高一些。使用 numexpr 提升在大型数据集上:

import numexpr as ne

def find_closest_v2(a, v):
    sidx = v.argsort()
    v_s = v[sidx]
    idx = np.searchsorted(v_s, a)
    idx[idx==len(v)] = len(v)-1
    idx0 = (idx-1).clip(min=0)
    
    p1 = v_s[idx]
    p2 = v_s[idx0]
    m = ne.evaluate('(idx!=0) & (abs(a-p1) >= abs(p2-a))', {'p1':p1, 'p2':p2, 'idx':idx})
    idx[m] -= 1
    out = sidx[idx]
    return out

计时
设置 :

N,M = 500,100000
a = np.random.rand(N,M)
v = np.random.rand(N)

In [22]: %timeit find_closest_v2(a, v)
4.35 s ± 21.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

In [23]: %timeit find_closest(a, v)
4.69 s ± 173 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)