快速且不占用内存的 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

你想拥有自己的实现吗？如果是这样，你可以使用 k-d tree within KNN, it's much more efficient, otherwise, you could use KNN library support GPU such knn_cuda

更新

你可以试试，cuml。

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

Fast and not memory expensive k nearest neighbours search

python

search

knn

更新