内存受限的编程竞赛：第K小珊瑚值求解方案问询

任务描述 #

给定T组测试用例，每组生成N个珊瑚的大小序列：

• 第一个珊瑚大小 ( S_1 = A )
• 后续珊瑚大小按公式 ( S_i = (B \times S_{i-1} + C) % D ) 计算（( i>1 )）
  需要找出所有珊瑚大小升序排列后的第K个元素。

约束条件 #

• 测试用例数量：( 1 \leq T \leq 3 )
• 序列长度与目标位置：( 1 \leq K \leq N \leq 10^7 )
• 参数范围：( 0 \leq A < D \leq 10^{18} )，( 1 \leq C, B \times D \leq 10^{18} )
• 资源限制：可用内存64MB，时间限制1.9秒

初始方案的内存问题 #

若直接存储所有( 10^7 )个8字节元素，总内存约76MB，超出64MB的限制。如果内存允许（≥80MB），可以使用以下直接生成+std::nth_element的方案：

#include <iostream>
#include <vector>
#include <iterator>
#include <algorithm>

using biggie = long long;

int main() {
    int t;
    std::cin >> t;
    int i, n, k, j;
    biggie a, b, c, d;
    std::vector<biggie>::iterator it_ans;
    for (i = 0; i != t; ++i) {
        std::cin >> n >> k >> a >> b >> c >> d;
        std::vector<biggie> lut{ a };
        lut.reserve(n);
        for (j = 1; j != n; ++j) {
            lut.emplace_back((b * lut.back() + c) % d);
        }
        it_ans = std::next(lut.begin(), k - 1);
        std::nth_element(lut.begin(), it_ans, lut.end());
        std::cout << *it_ans << '\n';
    }
    return 0;
}

核心问题 #

1. 在给定的内存与时间约束下，如何正确完成任务？
2. 无法存储所有元素时，能否用std::nth_element结合滑动窗口技术求解？

最优解决方案（参考代码） #

以下是满足资源限制的高效实现，利用优先队列（堆）来控制内存使用：

#include <iostream>
#include <queue>

using biggie = long long;

int main() {
    int t;
    std::cin >> t;
    int i, n, k, j, j_lim;
    biggie a, b, c, d, prev, curr;
    for (i = 0; i != t; ++i) {
        std::cin >> n >> k >> a >> b >> c >> d;
        if (k < n - k + 1) {
            std::priority_queue<biggie, std::vector<biggie>, std::less<biggie>> q;
            q.push(a);
            prev = a;
            for (j = 1; j != k; ++j) {
                curr = (b * prev + c) % d;
                q.push(curr);
                prev = curr;
            }
            for (; j != n; ++j) {
                curr = (b * prev + c) % d;
                if (curr < q.top()) {
                    q.pop();
                    q.push(curr);
                }
                prev = curr;
            }
            std::cout << q.top() << '\n';
        }
        else {
            std::priority_queue<biggie, std::vector<biggie>, std::greater<biggie>> q;
            q.push(a);
            prev = a;
            for (j = 1, j_lim = n - k + 1; j != j_lim; ++j) {
                curr = (b * prev + c) % d;
                q.push(curr);
                prev = curr;
            }
            for (; j != n; ++j) {
                curr = (b * prev + c) % d;
                if (curr > q.top()) {
                    q.pop();
                    q.push(curr);
                }
                prev = curr;
            }
            std::cout << q.top() << '\n';
        }
    }
    return 0;
}

方案解析 #

1. 内存优化：仅维护一个大小为( \min(K, N-K+1) )的堆，内存占用远低于64MB（即使K取5e6，堆的内存也仅约40MB，满足限制）。
2. 时间效率：每个元素入堆/出堆操作的时间复杂度为( \log M )（M为堆的大小），总时间复杂度为( O(N \log M) )，对于( N=10^7 )来说，完全符合1.9秒的时间限制。
3. 逻辑说明：
   • 当K较小时：用大顶堆保存当前最小的K个元素，堆顶就是第K小的元素。遍历后续元素时，若元素比堆顶小，则替换堆顶，保证堆始终保存最小的K个元素。
   • 当K较大时：等价于求第( N-K+1 )大的元素，用小顶堆保存当前最大的( N-K+1 )个元素，堆顶就是目标值。遍历后续元素时，若元素比堆顶大，则替换堆顶。

关于std::nth_element的可行性 #

std::nth_element需要直接访问整个元素集合来进行分区操作，而滑动窗口只能维护局部元素，无法覆盖全局范围，无法保证能找到全局第K小的值。因此无法用std::nth_element结合滑动窗口解决此问题。

内容的提问来源于stack exchange，提问作者Petar Ivanov