python - 加快列表之间 float 的比较

Question

我有一段代码执行以下操作:

• 从列表 b_lst 中获取一个 float ，索引 indx
• 检查此 float 是否位于索引 i 的 float 和列表 a_lst 中的下一个(索引 i+1)之间
• 如果是，则将 indx 存储在第三个列表 (c_lst) 的子列表中，其中该子列表的索引是左侧列表的索引 float 在 a_lst 中(即:i)
• 对 b_lst 中的所有 float 重复上述操作

这是一个 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

这比原始函数快约 130 倍。

---

加3

正在关注 jtaylor的建议我将 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

这比原始函数快约 470 倍。

Answer 1

你想看看 numpy 的 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])