Python/Numpy-быстро найти индекс в массиве, ближайшем к некоторому значению

У меня есть массив значений t, который всегда находится в порядке возрастания (но не всегда равномерно). У меня есть еще одно значение, x. Мне нужно найти индекс в t такой, что t[index] ближе всего к x. Функция должна возвращать ноль для x t.Макс().

Я написал две функции для этого. Первый, 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)

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

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, это мая будет медленнее, чем первый... (Не сравнивается)


использовать searchsorted:

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

print t.searchsorted(x)

Edit:

Ах да, я вижу, что это то, что вы делаете в 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