在字典中查找整数最近邻

我有一个dict需要整数键：

a = {}
a[1] = 100
a[55] = 101
a[127] = 102

我希望在询问时能够选择最近的邻居：

a[20] # should return a[1] = 100
a[58] # should return a[55] = 101
a[167] # should return a[127] = 102

有没有一种Python式的方法可以做到这一点？（我想这可以通过循环所有字典来完成，但这可能不是最优雅的解决方案？）

双索引同样的问题（也是整数）：

 b[90, 1] = 100, b[90, 55] = 101, b[90, 127] = 102
 b[70, 1] = 40, b[70, 45] = 41, b[70, 107] = 42

我希望能够得到 b[73, 40] = b[70, 45] = 41, 即 2D 平面中的最近邻。

Update：在对这个答案中的两种方法进行基准测试之后，second方法明显更好，以至于它应该almost严格优选。

下面的方法处理n- 尺寸相同：

class NearestDict(dict):
    def __init__(self, ndims):
        super(NearestDict, self).__init__()
        self.ndims = ndims

    # Enforce dimensionality
    def __setitem__(self, key, val):
        if not isinstance(key, tuple): key = (key,)
        if len(key) != self.ndims: raise KeyError("key must be %d dimensions" % self.ndims)
        super(NearestDict, self).__setitem__(key, val)

    @staticmethod
    def __dist(ka, kb):
        assert len(ka) == len(kb)
        return sum((ea-eb)**2 for (ea, eb) in zip(ka, kb))

    # Helper method and might be of use
    def nearest_key(self, key):
        if not isinstance(key, tuple): key = (key,)
        nk = min((k for k in self), key=lambda k: NearestDict.__dist(key, k))
        return nk

    def __missing__(self, key):
        if not isinstance(key, tuple): key = (key,)
        if len(key) != self.ndims: raise KeyError("key must be %d dimensions" % self.ndims)
        return self[self.nearest_key(key)]

Demo:

a = NearestDict(1)
a[1] = 100
a[55] = 101
a[127] = 102
print a[20]    # 100
print a[58]    # 100
print a[167]   # 102
print a.nearest_key(20)     # (1,)
print a.nearest_key(58)     # (55,)
print a.nearest_key(127)    # (127,)

b = NearestDict(2)
b[90, 1]   = 100
b[90, 55]  = 101
b[90, 127] = 102
b[70, 1]   = 40
b[70, 45]  = 41
b[70, 107] = 42
print b[73, 40] # 41
print b.nearest_key((73,40)) # (70, 45)

请注意，如果键存在，则查找不会比标准字典查找慢。如果钥匙does not存在，您计算每个现有键之间的距离。没有缓存任何内容，尽管我想您可以添加它。

Edit:

从建议的方法卡斯拉的回答 https://stackoverflow.com/a/29094947/736937以下方法使用与上面相同的类来实现scipy's cKDTree http://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.cKDTree.html:

请注意，还有一个额外的可选参数，regenOnAdd这将允许您推迟（重新）构建 KDTree，直到完成（大部分）插入之后：

from scipy.spatial import cKDTree

class KDDict(dict):
    def __init__(self, ndims, regenOnAdd=False):
        super(KDDict, self).__init__()
        self.ndims = ndims
        self.regenOnAdd = regenOnAdd
        self.__keys = []
        self.__tree = None
        self.__stale = False

    # Enforce dimensionality
    def __setitem__(self, key, val):
        if not isinstance(key, tuple): key = (key,)
        if len(key) != self.ndims: raise KeyError("key must be %d dimensions" % self.ndims)
        self.__keys.append(key)
        self.__stale = True
        if self.regenOnAdd: self.regenTree()
        super(KDDict, self).__setitem__(key, val)

    def regenTree(self):
        self.__tree = cKDTree(self.__keys)
        self.__stale = False

    # Helper method and might be of use
    def nearest_key(self, key):
        if not isinstance(key, tuple): key = (key,)
        if self.__stale: self.regenTree()
        _, idx = self.__tree.query(key, 1)
        return self.__keys[idx]

    def __missing__(self, key):
        if not isinstance(key, tuple): key = (key,)
        if len(key) != self.ndims: raise KeyError("key must be %d dimensions" % self.ndims)
        return self[self.nearest_key(key)]

输出与上面的方法相同。

基准测试结果

了解三种方法的性能（NearestDict, KDDict(True)（插入时再生），以及KDDict(False)（延迟再生）），我简要地对它们进行了基准测试。

我进行了 3 次不同的测试。在测试中保持不变的参数是：

• 测试迭代次数：每个测试我都做了5次，并且花费了最短的时间。 （笔记timeit.repeat默认为 3)。
• 点边界：我生成了 0
• 查找次数：我分别对插入和查找进行计时。下面的三个测试均使用 10,000 次查找。

第一个测试使用 4 个维度的按键和 1,000 次插入。

{'NDIMS': 4, 'NITER': 5, 'NELEMS': 1000, 'NFINDS': 10000, 'DIM_LB': 0, 'DIM_UB': 1000, 'SCORE_MUL': 100}
insert::NearestDict       0.125
insert::KDDict(regen)    35.957
insert::KDDict(defer)     0.174
search::NearestDict    2636.965
search::KDDict(regen)    49.965
search::KDDict(defer)    51.880
  

第二次测试使用 4 维按键和 100 次插入。我想改变插入的数量，看看随着字典密度的变化，这两种方法的表现如何。

{'NDIMS': 4, 'NITER': 5, 'NELEMS': 100, 'NFINDS': 10000, 'DIM_LB': 0, 'DIM_UB': 1000, 'SCORE_MUL': 100}
insert::NearestDict       0.013
insert::KDDict(regen)     0.629
insert::KDDict(defer)     0.018
search::NearestDict     247.920
search::KDDict(regen)    44.523
search::KDDict(defer)    44.718
  

第三次测试使用了 100 次插入（与第二次测试类似），但使用了 12 个维度。我想看看随着关键维度的增加，这些方法的表现如何。

{'NDIMS': 12, 'NITER': 5, 'NELEMS': 100, 'NFINDS': 10000, 'DIM_LB': 0, 'DIM_UB': 1000, 'SCORE_MUL': 100}
insert::NearestDict       0.013
insert::KDDict(regen)     0.722
insert::KDDict(defer)     0.017
search::NearestDict     405.092
search::KDDict(regen)    49.046
search::KDDict(defer)    50.601
  

讨论

KDDict连续再生（KDDict(True)) 要么稍微快一些（在查找中）或者相当较慢（插入时）。因此，我将其排除在讨论之外，并将重点放在NearestDict and KDDict(False)，现在简称为KDDict

结果令人惊讶地支持具有延迟再生的 KDDict。

对于插入，在所有情况下，KDDict 的表现都比 NearestDict 稍差。这是预料之中的，因为有额外的列表追加操作。

对于搜索，在所有情况下，KDDict 都执行明显更好比最近的字典。

随着字典稀疏度的降低/密度的增加，NearestDict 的性能下降程度比 KDDict 更大。当从 100 个键增加到 1000 个键时，NearestDict 搜索时间增加了 9.64 倍，而 KDDict 搜索时间仅增加了 0.16 倍。

随着字典维数的增加，NearestDict 的性能下降幅度大于 KDDict。当从 4 维变为 12 维时，NearestDict 搜索时间增加了 0.64 倍，而 KDDict 搜索时间仅增加了 0.13 倍。

鉴于此，以及两个类的相对相同的复杂性，如果您可以访问 scipy 工具包，请使用KDDict强烈推荐这种方法。