python - TLE 计算列表中指定范围内的元素数量

Question

有一个未排序的列表 a 和一个范围列表，如 ranges = [(10, 20), (30, 50), (15, 35) ...]。 a 中的最大值为 uint64_t。目标是计算每个范围的元素数量。正常的解决方案非常直观，只需统计范围内的元素并打印结果即可。但是问题来自在线法官。我厌倦了几个解决方案，但对于每个解决方案，OJ 都给出了超出时间限制。

a 的最大长度为 10,000,000，ranges 的最大长度为 1,000,000。

具有 1000 万个随机数的测试列表 a 和具有 100 万对范围的 ranges:

import numpy as np

a = list(np.random.randint(low=1, high=0x7fffffffffffffff, size=10_000_000))

ranges = []
for _ in range(1_000_000):
    x, y = np.random.randint(low=1, high=0x7fffffffffffffff, size=2)
    ranges.append((x, y) if x < y else (y, x))

第一种解决方案是:

import bisect

a.sort()

low_d = {}
up_d = {}

def count(r):
    low, up = r

    if low not in low_d:
        l = bisect.bisect_left(a, low)
        low_d[low] = l
    else:
        l = low_d[low]

    if up not in up_d:
        u = bisect.bisect_right(a, up, lo=l)
        up_d[up] = u
    else:
        u = up_d[up]

    return u - l

result = [*map(count, ranges)]

缺点很明显，当a非常大时，sort()非常耗时。

The original second solution is much slower than the above solution.

Abandoned.

两种解决方案都会导致 TLE 错误。我使用的OJ就像一个黑盒子，我不知道它用来测试程序的测试例子。

由于程序运行在OJ上，所以不允许使用numpy。

有什么方法可以优化性能吗？

Answer 1

我设法将机器上的时间从 13.1 秒减少到 11.2 秒

我的最终代码:

from bisect import bisect_left, bisect_right
def f0_4(a, ranges, n_pre_b):
    a.sort()
    blen = [x*len(a)//n_pre_b for x in range(n_pre_b)]
    b1 = [a[i] for i in blen]
    blen.append(len(a))
    b1.append(a[-1])
    res = []
    for low, up in ranges: 
        low_pre_b_i = bisect_left(b1,low)
        lo = blen[low_pre_b_i-1]
        hi = blen[low_pre_b_i]
        l = bisect_left(a, low, lo=lo, hi=hi)
        high_pre_b_i = bisect_left(b1,up)
        lo = blen[high_pre_b_i-1]
        hi = blen[high_pre_b_i]
        if l > lo:
            res.append(bisect_right(a, up, lo=l, hi=hi)-l)
        else:
            res.append(bisect_right(a, up, lo=lo, hi=hi)-l)

    return res
res = f0_4(a,ranges,16384)

什么和为什么:

1. 我删除了函数调用“count”，因为在 python 中每次调用都是开销
2. 我删除了 bisect 值的缓存，因为具有相同 bisect 值的概率很小。测试表明没有缓存速度更快
3. 我为 a 列表预先计算了一些范围。 16384 个预先计算的值是最佳的。这极大地提高了速度
4. 我通过更改导入将 bisect.bisect_left 替换为 bisect_left。 bisect_right 也一样。点调用在 python 中有开销

我本来应该做的，但这是违反规则的(如果我错了请纠正我。下面的方法可以将速度提高 100 倍):

1. 使用 numpy 数组而不是 python 列表。排序速度提高了 10 倍。允许使用快速 numba 和 cython 代码。元素访问速度更快
2. 在 cython 或 numba 中对均匀分布的整数使用特殊的排序算法。可能会比原来的排序提高 100 倍。这种排序与数组大小成线性关系。这对于 python 列表和 numpy 数组使用的一般排序算法是不可能的
3. 使用 numpy.searchsorted() 而不是 bisect
4. 将所有 for 循环转换为 numba 或 cython 代码。 python 中的循环非常低效

通过实现上述所有内容，我可能会编写速度提高大约 100 倍的 C++ 代码。在我的 C++ 小经验中，它总是比 cython 或 numpy 的任何优化都快。但 C++ 是一个不同的问题。

我试过但更糟的是:

1. 对范围进行排序而不对 a 进行排序。与提问者的版本相比，时间增加了 4 倍。很可能是由于代码的复杂性要高得多
2. 用数字除以最大数来猜测二分法。然后在附近用二分法搜索。比提问者的版本快，但比具有预先计算间隔的最终版本慢。
3. 使用字典而不是列表。更糟

代码的行分析器:

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     1                                            def f0_4(a, ranges, n_pre_b):
     2         1    6465309.0 6465309.0     39.8      a.sort()
     3         1       3088.0    3088.0      0.0      blen = [x*len(a)//n_pre_b for x in range(n_pre_b)]
     4         1       4661.0    4661.0      0.0      b1 = [a[i] for i in blen]
     5         1          4.0       4.0      0.0      blen.append(len(a))
     6         1          1.0       1.0      0.0      b1.append(a[-1])
     7         1          1.0       1.0      0.0      res = []
     8   1000001     540421.0       0.5      3.3      for low, up in ranges: 
     9   1000000    1180737.0       1.2      7.3          low_pre_b_i = bisect.bisect_left(b1,low)
    10   1000000     608838.0       0.6      3.7          lo = blen[low_pre_b_i-1]
    11   1000000     490782.0       0.5      3.0          hi = blen[low_pre_b_i]
    12   1000000    2064953.0       2.1     12.7          l = bisect.bisect_left(a, low, lo=lo, hi=hi)
    13   1000000    1212568.0       1.2      7.5          high_pre_b_i = bisect.bisect_left(b1,up)
    14   1000000     606433.0       0.6      3.7          lo = blen[high_pre_b_i-1]
    15   1000000     492544.0       0.5      3.0          hi = blen[high_pre_b_i]
    16   1000000     460683.0       0.5      2.8          if l > lo:
    17        54        103.0       1.9      0.0              res.append(bisect.bisect_right(a, up, lo=l, hi=hi)-l)
    18                                                   else:
    19    999946    2132459.0       2.1     13.1              res.append(bisect.bisect_right(a, up, lo=lo, hi=hi)-l)
    20                                                   
    21         1          1.0       1.0      0.0      return res