我需要一种包装我的KDTree的方法，以在树的另一侧找到最近的邻居

如何解决我需要一种包装我的KDTree的方法，以在树的另一侧找到最近的邻居

找到cKDTree可以找到最近的邻居，我一直在尝试运行该程序并获得正确输出所花费的时间。

我发现树只匹配一行，并且不检查最近的邻居是否在树的另一侧，因此输出不正确。

需要找到一种包装树以获取正确输出的方法。

# Reference code that works with brute force method

import numpy as np
from multiprocessing.dummy import Pool as ThreadPool
import time
import errno

N = np.int(input('N = '))
seed = np.int(input('Random number seed = '))
np.random.seed(seed)
coords = np.random.random((3,N)).transpose()
start_time = time.time()


def NN(point):
    dist = np.abs(np.subtract(coords,point))
    dist = np.where(dist > 0.5,dist - 1,dist)
    return (np.square(dist)).sum(axis=-1).argsort()[1]


pool = ThreadPool(12)
match = pool.map(NN,coords)
pool.close()
pool.join()

end_time = time.time()
print('Elapsed time = ',repr(end_time - start_time))

# generate filename from N and seed
filename = 'pyneigh' + str(N) + '_' + str(seed)
# if a file with this name already exists read in the nearest neighbour
# list found before and compare this to the current list of neighbours,# else save the neighbour list to disk for future reference
try:
    fid = open(filename,'rb')
    matchold = np.loadtxt(fid)
    fid.close()
    if (matchold == match).all():
        print('Checked match')
    else:
        print('Failed match')
except OSError as e:
    if e.errno == errno.ENOENT:
        print('Saving neighbour list to disk')
        fid = open(filename,'wb')
        np.savetxt(fid,match,fmt="%8i")
        fid.close()
    else:
        raise
#

#########################################################################################################
# My current code

import numpy as np
import time
import errno
from scipy.spatial import cKDTree

N = np.int(input('N = '))
seed = np.int(input('Random number seed = '))
np.random.seed(seed)
# set up random 3-d positions
pos = np.random.random((3,N))
start_time = time.time()

# You may only change the code in this cell (i.e. between here and end_time)
#
pos2 = pos.T

tree = cKDTree(pos2)
nearest_dist,nearest_ind = tree.query(pos2,k=3)  # k=3 nearest neighbors where k1 = identity

match = nearest_ind[:,1].T

end_time = time.time()
print('Elapsed time = ',fmt="%8i")
        fid.close()
    else:
        raise

#########################################################################################################