快速且不占用内存的 k 最近邻搜索

Question

我正在尝试为不同数据集中的新点数组中的每个元素找到最近的邻居，这会很快而且不会占用大量内存。我更关心的是为更多邻居而不是更多维度调整代码。

基于https://glowingpython.blogspot.com/2012/04/k-nearest-neighbor-search.html?showComment=1355311029556#c8236097544823362777 我写了 k 最近邻搜索，但它的内存非常广泛。在我的实际问题中，我有 100 万个值要搜索，100k 个点需要匹配，100 万 x 10k 数组估计为 600GiB。

有没有更好的办法？

我曾尝试使用 bisect（基于 from list of integers, get number closest to a given value），但我必须循环 100k 次，这需要一些时间，尤其是我要进行多次搜索。

小型数据集的好代码 - 能够找到 K 个最近的邻居，并且可以轻松地添加到许多维度（按维度循环）：

def knn_search(search_for, search_in, K = 1, 
               return_col = ["ID"],
               col = 'A'):
        
    
    #print(col)
    a_search_in  = array(search_in[col])
    a_search_for = array(search_for[col])
    
    #print('a')
    a = np.tile(a_search_for, [a_search_in.shape[0], 1]).T
    #print('b')
    b = np.tile(a_search_in,  [a_search_for.shape[0], 1])
    #print('tdif')
    t_diff =  a - b
        
    #print('suma')
    diff = np.square(t_diff)

    # sorting
    idx  = argsort(diff)
    
    
    # return the indexes of K nearest neighbours
    if search_for.shape[0] == 1:
        return idx[:K]
    elif K == 1:
        return search_in.iloc[np.concatenate(idx[:,:K]), :][return_col]
    else:
        tmp = pd.DataFrame()
        for i in range(min(K, search_in.shape[0])):
            tmp = pd.concat([tmp.reset_index(drop=True), 
                             search_in.iloc[idx[:,i], :][[return_col]].reset_index(drop=True)], 
                            axis=1)
        return tmp

1 维和 1 个邻居的好代码：

def knn_search_1K_1D(search_for, search_in, 
           return_col = ["ID"],
           col = 'A'):
    sort_search_in = search_in.sort_values(col).reset_index()
        idx = np.searchsorted(sort_search_in[col], search_for[col])
        idx_pop = np.where(idx > len(sort_search_in) - 1, len(sort_search_in) - 1, idx)
    
    t = sort_search_in.iloc[idx_pop  , :][[return_col]]
    search_for_nn = pd.concat([search_for.add_prefix('').reset_index(drop=True), 
                             t.add_prefix('nn_').reset_index(drop=True)], 
                            axis=1)

K 最近邻 > 1 维和 1 维的当前工作解决方案，但在上述实际情况下需要一个多小时来计算

def knn_search_nK_1D(search_for, search_in, K = 1, 
               return_col = ["ID"],
               col = 'A'):
    t = []
    #looping one point by one 
    for i in range(search_for.shape[0]):
        y = search_in[col]
        x = search_for.iloc[i, :][col]
        nn = np.nanmean(search_in.iloc[np.argsort(np.abs(np.subtract(y, x)))[0:K], :][return_col])
        t.append(nn)
    search_for_nn = search_for
    search_for_nn['nn_' + return_col] = t

示例数据：

search_for = pd.DataFrame({'ID': ["F", "G"],
                          'A' : [-1,  9]})

search_in = pd.DataFrame({'ID': ["A", "B", "C", "D", "E"],
                          'A' : [1,    2,   3,   4,   5 ]})



t = knn_search(search_for = search_for , 
               search_in  = search_in,
               K = 1, 
               return_col = ['ID'],
               col = 'A')
print(t)
#  ID
#0  A
#4  E

Answer 1

您想拥有自己的实现吗？如果是这样，您可以在KNN 中使用k-d tree，这样效率更高，否则，您可以使用KNN 库支持GPU，例如knn_cuda

---

更新

你可以试试，cuml。