python - 当坐标保存在对象中时，使用 python 中的 kd-tree 查找 k 个最近邻

Question

我需要为集合中的每个对象找到k个最近邻居。每个对象都有其坐标作为属性。为了解决这个任务，我尝试使用 scipy 中的 spatial.KDTree 。如果我使用列表或元组来表示一个点，它工作得很好，但它不适用于对象。我在类中实现了 __getitem__ 和 __len__ 方法，但是 KDTree 实现要求我的对象提供不存在的坐标轴(比如第三个坐标)二维点)。

这是一个重现问题的简单脚本:

from scipy import spatial

class Unit:

    def __init__(self, x,y):
        self.x = x
        self.y = y


    def __getitem__(self, index):        
        if index == 0:
            return self.x
        elif index == 1:
            return self.y
        else:          
            raise Exception('Unit coordinates are 2 dimensional')


    def __len__(self):        
        return 2



#points = [(1, 1), (2, 2), (3, 3), (4, 4), (5, 5)]
#points = [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]]
points = [Unit(1,1), Unit(2,2), Unit(3,3), Unit(4,4), Unit(5,5)]

tree = spatial.KDTree(points)

#result = tree.query((6,6), 3)
result = tree.query(Unit(6,6), 3)

print(result)

我没有必要使用这个特定的实现或库甚至算法，但要求是处理对象。

附注我可以向每个对象添加 id 字段，并将所有坐标移动到单独的数组中，其中索引是对象 id 。但如果可能的话，我仍然想避免这种方法。

Answer 1

docs对于 scipy.spatial.KDTree声明data参数应该是array_like这通常意味着“可转换为 numpy 数组”。事实上，初始化的第一行尝试将数据转换为 numpy 数组，正如您在 source code 中看到的那样。 :

class KDTree(object):
    """ ... """
    def __init__(self, data, leafsize=10):
        self.data = np.asarray(data)

所以你想要实现的是一个对象，以便它们的列表可以很好地转换为 numpy 数组。这是hard to exactly define因为 numpy 尝试了多种方法将对象放入数组中。然而，包含许多相同长度序列的迭代绝对符合资格。

您的Unit对象基本上是一个序列，因为它实现了 __len__和__getitem__以及从 0 开始的连续整数的索引。Python 知道你的序列何时结束，因为它命中 IndexError 。但是你的__getitem__提出 Exception相反，在一个坏索引上。因此，从这两种方法提供顺序迭代的正常机制被打破。相反，请提出 IndexError ，你会很好地转换:

class Unit:
    def __init__(self, x, y):
        self.x = x
        self.y = y

    def __getitem__(self, index):        
        if index == 0:
            return self.x
        elif index == 1:
            return self.y
        raise IndexError('Unit coordinates are 2 dimensional')

    def __len__(self):        
        return 2

现在我们可以毫无问题地检查这些转换为 numpy 数组的列表:

In [5]: np.array([Unit(1, 1), Unit(2, 2), Unit(3, 3), Unit(4, 4), Unit(5, 5)])
Out[5]:
array([[1, 1],
       [2, 2],
       [3, 3],
       [4, 4],
       [5, 5]])

所以，我们初始化 KDTree 应该没有问题现在。这就是为什么如果您将坐标存储在内部列表中并只是推迟 __getitem__到该列表，或者只是将您的坐标视为一些简单的序列(例如列表或元组)，就可以了。

使用像这样的简单类的更简单的方法是使用 namedtuples或类似的，但对于更复杂的对象，将它们转换成序列是一个好方法。