快速且不消耗内存的k个最近邻居搜索

问题描述

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

基于https://glowingpython.blogspot.com/2012/04/k-nearest-neighbor-search.html?showComment=1355311029556#c8236097544823362777 我已经写了k最近邻居搜索，但是它占用大量内存。在我的实际问题中，我需要搜索100万个值，并且需要匹配10万个点，对于100万x 1万个数组，估计为600GiB。

有更好的方法吗？

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

适用于小型数据集的良好代码-能够找到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,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,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,return_col = ['ID'],col = 'A')
print(t)
#  ID
#0  A
#4  E

解决方法

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

---

更新

您可以尝试，cuml。