快速找到数组中最接近某个值的索引

Question

我有一个值的数组，叫做 t，它的值是一直在增加的（但不是均匀间隔的）。我还有一个单独的值，叫做 x。我需要找到 t 中的一个索引，使得 t[index] 和 x 的差距最小。如果 x 小于 t 中的最小值，就返回 0；如果 x 大于 t 中的最大值，就返回最大索引（或者 -1）。

我写了两个函数来实现这个功能。第一个函数 f1 在这个简单的时间测试中快得多。但我也喜欢第二个函数，因为它只有一行代码。这个计算会在一个很大的数组上进行，可能每秒会执行很多次。

有没有人能想出一个其他的函数，它的速度和第一个差不多，但代码看起来更简洁？或者有没有比第一个更快的（速度是最重要的）？

谢谢！

代码：

import numpy as np
import timeit

t = np.arange(10,100000)         # Not always uniform, but in increasing order
x = np.random.uniform(10,100000) # Some value to find within t

def f1(t, x):
   ind = np.searchsorted(t, x)   # Get index to preserve order
   ind = min(len(t)-1, ind)      # In case x > max(t)
   ind = max(1, ind)             # In case x < min(t)
   if x < (t[ind-1] + t[ind]) / 2.0:   # Closer to the smaller number
      ind = ind-1
   return ind

def f2(t, x):
   return np.abs(t-x).argmin()

print t,           '\n', x,           '\n'
print f1(t, x),    '\n', f2(t, x),    '\n'
print t[f1(t, x)], '\n', t[f2(t, x)], '\n'

runs = 1000
time = timeit.Timer('f1(t, x)', 'from __main__ import f1, t, x')
print round(time.timeit(runs), 6)

time = timeit.Timer('f2(t, x)', 'from __main__ import f2, t, x')
print round(time.timeit(runs), 6)

Answer 1

np.searchsorted 是一种二分查找的方法，它每次都会把数组分成两半。所以你需要把它实现成返回小于 x 的最后一个值，而不是返回零。

看看这个算法（来自 这里）：

def binary_search(a, x):
    lo=0
    hi = len(a)
    while lo < hi:
        mid = (lo+hi)//2
        midval = a[mid]
        if midval < x:
            lo = mid+1
        elif midval > x: 
            hi = mid
        else:
            return mid
    return lo-1 if lo > 0 else 0

只需要把最后一行替换掉（原来是 return -1）。同时也改了一下参数。

因为这些循环是用 Python 写的，所以它的速度可能会比第一个算法慢……（没有进行性能测试）

Answer 2

使用searchsorted：

t = np.arange(10,100000)         # Not always uniform, but in increasing order
x = np.random.uniform(10,100000)

print t.searchsorted(x)

编辑：

哦，对了，我明白你在f1里是怎么做的。也许下面的f3比f1更容易理解。

def f3(t, x):
    ind = t.searchsorted(x)
    if ind == len(t):
        return ind - 1 # x > max(t)
    elif ind == 0:
        return 0
    before = ind-1
    if x-t[before] < t[ind]-x:
        ind -= 1
    return ind

Answer 3

这看起来快多了（对我来说，使用的是Python 3.2-win32，numpy 1.6.0）：

from bisect import bisect_left
def f3(t, x):
    i = bisect_left(t, x)
    if t[i] - x > 0.5:
        i-=1
    return i

---

输出结果：

[   10    11    12 ..., 99997 99998 99999]
37854.22200356027
37844
37844
37844
37854
37854
37854
f1 0.332725
f2 1.387974
f3 0.085864