Ускорить сравнение поплавков между списками

Question

У меня есть блок кода, который делает следующее:

• взять поплавок из списка, b_lst ниже индекса indx
• чека, если этот поплавок находится между поплавок индекса i, а следующий (индекс i+1) в списке a_lst
• , если он есть, то хранить indx в подсписке третьего списка (c_lst), где индекс этого подсписка является индекс левого поплавка в a_lst (т.е. i)
• повторить для всех поплавков в b_lst

Вот MWE, который показывает, что делает этот код:

import numpy as np 
import timeit 

def random_data(N): 
    # Generate some random data. 
    return np.random.uniform(0., 10., N).tolist() 

# Data lists. 
# Note that a_lst is sorted. 
a_lst = np.sort(random_data(1000)) 
b_lst = random_data(5000) 
# Fixed index value (int) 
c = 25 

def func(): 
    # Create empty list with as many sub-lists as elements present 
    # in a_lst beyond the 'c' index. 
    c_lst = [[] for _ in range(len(a_lst[c:])-1)] 

    # For each element in b_lst. 
    for indx,elem in enumerate(b_lst): 

     # For elements in a_lst beyond the 'c' index. 
     for i in range(len(a_lst[c:])-1): 

      # Check if 'elem' is between this a_lst element 
      # and the next. 
      if a_lst[c+i] < elem <= a_lst[c+(i+1)]: 

       # If it is then store the index of 'elem' ('indx') 
       # in the 'i' sub-list of c_lst. 
       c_lst[i].append(indx) 

    return c_lst 

print func() 
# time function. 
func_time = timeit.timeit(func, number=10) 
print func_time 

Этот код работает так, как должно, но мне действительно нужно, чтобы улучшить его производительность, так как это замедляет остальную часть моего кода.

---

Добавить

Это оптимизированная функция, основанная на принятом ответе. Это довольно уродливо, но он выполняет свою работу.

def func_opt(): 
    c_lst = [[] for _ in range(len(a_lst[c:])-1)] 
    c_opt = np.searchsorted(a_lst[c:], b_lst, side='left') 
    for elem in c_opt: 
     if 0<elem<len(a_lst[c:]): 
      c_lst[elem-1] = np.where(c_opt==elem)[0].tolist() 
    return c_lst 

В моих тестах это ~ 7 раз быстрее, чем исходная функция.

---

Добавить 2

Гораздо быстрее, не используя np.where:

def func_opt2(): 
    c_lst = [[] for _ in range(len(a_lst[c:])-1)] 
    c_opt = np.searchsorted(a_lst[c:], b_lst, side='left') 
    for indx,elem in enumerate(c_opt): 
     if 0<elem<len(a_lst[c:]): 
      c_lst[elem-1].append(indx) 
    return c_lst 

Это ~ 130x быстрее, чем исходная функция.

---

Добавить 3

Следуя совету jtaylor «s я преобразованного результат np.searchsorted к списку с .tolist():

def func_opt3(): 
    c_lst = [[] for _ in range(len(a_lst[c:])-1)] 
    c_opt = np.searchsorted(a_lst[c:], b_lst, side='left').tolist() 
    for indx,elem in enumerate(c_opt): 
     if 0<elem<len(a_lst[c:]): 
      c_lst[elem-1].append(indx) 
    return c_lst 

Это ~ 470x быстрее, чем исходная функция.

Answer 1

Вы хотите взглянуть на numpy's searchsorted. Вызов

np.searchsorted(a_lst, b_lst, side='right') 

возвращает массив индексов, такие же длины, как b_lst, держа перед каким элементом в a_lst они должны быть вставлены, чтобы сохранить порядок. Он будет очень быстрым, поскольку он использует бинарный поиск, и цикл происходит в C. Затем вы можете создать свои подмассивы с фантастическим индексированием, например.:

>>> a = np.arange(1, 10) 
>>> b = np.random.rand(100) * 10 
>>> c = np.searchsorted(a, b, side='right') 
>>> b[c == 0] 
array([ 0.54620226, 0.40043875, 0.62398925, 0.40097674, 0.58765603, 
     0.14045264, 0.16990249, 0.78264088, 0.51507254, 0.31808327, 
     0.03895417, 0.92130027]) 
>>> b[c == 1] 
array([ 1.34599709, 1.42645778, 1.13025996, 1.20096723, 1.75724448, 
     1.87447058, 1.23422399, 1.37807553, 1.64118058, 1.53740299])