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

Question

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

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

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

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() 非常耗时。

原来的第二种方案比上面的方案慢很多。

被遗弃了。

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

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

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

Answer 1

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

我的最终代码：

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

Answer 2

一个稍微快一点的方法是使用列表推导。为了在特定情况下加快速度（通常由在线评委测试），我还使用了set(ranges) 以防您的范围有重复，但如果您知道几乎没有重复，可以将其删除。

这里是代码位：

import random

# generate random data that looks like yours
a = [random.randint(0, 30) for i in range(1000)]
ranges = [(random.randint(0, 20), random.randint(10, 30)) for i in range(1000)]

# using dictionary (your code)
def old_count(a, ranges):

    d = {}

    for i in range(len(ranges)):
        d[i] = 0

    for v in a:
        for i, (l, u) in enumerate(ranges):
            if l <= v <= u:
                d[i] += 1

    return d.values()

# using list comprehension (new code)
def count(a, ranges):
    return [((i, j), sum([i<=x<=j for x in a])) for (i, j) in set(ranges)]

时间方面的结果在我的笔记本电脑上是这样的：

# using dictionary (your code)

# for 1000 items in a and 1000 ranges
%timeit old_count(a, ranges)
# > 155 ms ± 1.02 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# for 10000 items in a and 10000 ranges
%timeit count(a, ranges)
# > 19.4 s ± 1.91 s per loop (mean ± std. dev. of 7 runs, 1 loop each)

# using list comprehension (new code)

# for 1000 items in a and 1000 ranges
%timeit count(a, ranges)
# > 39.9 ms ± 1.39 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

# for 10000 items in a and 10000 ranges
%timeit count(a, ranges)
# > 593 ms ± 112 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

不过可能有更快/更好的解决方案。

有几点可以改进，我只列举几个，但可能还有其他的：

• a 中也可能有重复项
• 某些范围可能包含其他范围，因此无需检查某些范围
• 如果范围列表是有序的，则不值得检查所有范围

Answer 3

此 C++ 代码在使用 -O2 编译的 1.9 秒内运行，而我在此硬件上的最佳 Python 代码为 13.2 秒（与 Python 中的基准测试相比，这是一个较慢的硬件）。

可能的改进：

1. 上等分线应在下等分线上方搜索
2. 使用预先计算的二等分值，如在 python 代码中
3. 从Unisort: an Algorithm to Sort Uniformly Distributed Numbers in O(n) Time. R.T. Ionescu 2013 实现 Unisort 算法

代码：

#include <iostream>
#include <string>
#include <random>
#include <cstdint>
#include <array>
#include <chrono>
#include <algorithm>
#include <iterator>

int tdiff(std::chrono::time_point<std::chrono::system_clock> start, std::chrono::time_point<std::chrono::system_clock> _end) {
    int result;
    result = (std::chrono::duration_cast<std::chrono::milliseconds>(_end - start)).count();
    return result; 
}

int main() 
{
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<int64_t> dis(1, 0x7fffffffffffffff);

    #define A_SIZE 10000000
    #define R_SIZE 1000000
    std::vector<int64_t> a(A_SIZE);
    int a_size = A_SIZE;
    int r_size = R_SIZE;
    for (int i=0; i<a_size; i++){
        a[i] = dis(gen);
    }
    std::vector<std::vector<int64_t>> ranges1(R_SIZE, std::vector<int64_t>(2));

    int64_t x,y;
    for (int i=0; i<ranges1.size(); i++){
        x = dis(gen);
        y = dis(gen);
        if (x < y){
            ranges1[i] = {x,y};
        }else{
            ranges1[i] = {y,x};
        }
    }
    std::chrono::time_point<std::chrono::system_clock> start, _end;

    start = std::chrono::system_clock::now();
    std::sort(a.begin(), a.end());
    std::vector<int64_t> counts(A_SIZE);
    std::vector<int64_t>::iterator l;
    std::vector<int64_t>::iterator u;
    for (int i=0; i<r_size; i++){
        l = std::lower_bound(a.begin(),a.end(),ranges1[i][0]);
        u = std::upper_bound(a.begin(),a.end(),ranges1[i][1]);
        counts[i] = (int64_t)std::distance(l,u);
    }
    _end = std::chrono::system_clock::now();
    std::cout << tdiff(start, _end) << "\n";

    std::cout << counts[0] << "\n";
    return 0;
}