如何在 O(logn) 中找到 pandas 数据框列中与输入值 x 最接近的值?

How to find k nearest values in a pandas data frame column to an input value x in O(logn)?

我有一个包含两列的数据框:1) ID:代表样本 ID 的随机整数,2) A:浮点数

size_df1 = 1000
df1 = pd.DataFrame(np.random.random_sample((size_df1)), columns=list('A'))
df1['ID'] = random.sample(range(0, size_df1), size_df1)

给定像 x=0.21 这样的输入,如何在 df1['A'] 中找到 10 个(或任何其他整数,例如 k)最接近 x 的值,在 log(n) 中,其中 n 是行数df1。请注意,这应该在没有替换的情况下完成,每次我在 df1['A'] 中找到这 10 个最接近的值时,我应该删除这些值或以某种方式标记它们并且不将它们用于下一个 x。这完全可以在 logn 中解决吗?谢谢

.nsmallest()可以轻松找到k个最小的值,最接近的值是绝对差值最小的值:

>>> (df1['A'] - 0.21).abs().nsmallest(10)
969    0.000014
889    0.000442
779    0.003299
259    0.003637
843    0.003700
84     0.003818
651    0.004264
403    0.004360
648    0.004421
543    0.005088
Name: A, dtype: float64

如果您想访问匹配的行,您可以重用它的索引:

>>> df1.loc[(df1['A'] - 0.21).abs().nsmallest(10).index]
            A   ID
969  0.210014  237
889  0.210442  225
779  0.206701  127
259  0.213637  883
843  0.206300  330
84   0.206182   17
651  0.205736   64
403  0.205640  388
648  0.214421  964
543  0.204912  616

请注意 nsmallest 的文档说:

Faster than .sort_values().head(n) for small n relative to the size of the Series object.

关于复杂性的一个词,因为你的值没有排序:

  • 如果您想找到 1 个最接近的值,则最低复杂度是 O(n)
  • 您可以进行类似二分查找的操作来获得 O(log(n)),但这需要先进行排序 - 所以它实际上是 O(n log(n)).

假设您的数据框按 A:

排序
>>> df1.sort_values('A', inplace=True)

那我们可以尝试使用排序搜索功能,其中returns行号(不是索引值):

>>> df1['A'].searchsorted(0.21)
197

这意味着我们可以使用它来找到 k 最接近的候选者,然后在这个 2k 数据帧上使用我们之前的方法:

def find_closest(df, val, k):
    return df.loc[df['A'].sub(val).abs().nsmallest(k).index]

def find_closest_sorted(df, val, k):
    closest = df['A'].searchsorted(val)
    if closest < k:
        return find_closest(df.iloc[:closest + k], val, k)

    return find_closest(df.iloc[closest - k:closest + k], val, k)
>>> find_closest_sorted(df1, 0.21, 10)
            A   ID
969  0.210014  237
889  0.210442  225
779  0.206701  127
259  0.213637  883
843  0.206300  330
84   0.206182   17
651  0.205736   64
403  0.205640  388
648  0.214421  964
543  0.204912  616

复杂度应该在这里:

  • O(n log(n)) 用于排序(可以分摊到多次查找中)
  • O(log(n)) 用于排序搜索
  • O(k) 最后一步。