Codeforces 1837F Editorial for Two

题目来自 Codeforces Educational Round #149 (Div.2) F。link

  This article provides a feasible solution to a certain competitive programming problem.   If you have better solutions, feel free to share with me!

Problem Description

给出长度为 \(n\) 的数列 \(\{a_1,a_2,...,a_n\}\) \((1\leq a_i\leq10^9)\)，选出一个长度为 \(k\) 的子序列 \((1\leq k\leq n\leq 3\cdot 10^5)\)，并将该子序列分为左右两部分。找到合适的划分方案，使得 \(\{\max(\)左侧元素之和，右侧元素之和\()\}\) 最小，并求出该最小值。数据组数 \(1\leq t\leq 10^4\)，保证 \(\sum n\leq 3\times 10^5\)。

Tutorial

首先最小化最大值一眼二分答案。考虑当前二分的值为 \(x\)。问题转化为，在长度为 \(n\) 的数列中，选出两个长度之和为 \(k\) 的子序列，使得这两个子序列的元素之和均 \(\leq x\)。并且，这两个子序列不重合：第一个子序列的末尾元素一定在第二个子序列的起始元素之前。

考虑维护前缀数组 \(pre_i\) 与后缀数组 \(suf_i\)，标记在前缀 \(a_{1...i}\) 或后缀 \(a_{i...n}\) 中，最多能够选择多少个元素，使得这些元素之和 \(\leq x\)。若我们能够成功计算出 \(pre_i\) 与 \(suf_i\)，那么只要使得 \(pre_i+suf_{i+1}\geq k\) 的 \(i\) 存在，就能选出两个不重合，且元素之和 \(\leq x\) 的两个子序列。

关于 \(pre_i\) 与 \(suf_i\) 的计算，这里采取一个简单但巧妙的贪心算法：以 \(pre_i\) 为例，维护一个大根堆 (或优先队列)，堆内所有元素的和 \(sum\) 与堆的大小 \(size\)。对于新元素 \(a_i\)，先将 \(a_i\) 加入堆中，并更新 \(sum\), \(size\)。若 \(sum\leq x\)，\(pre_i=size\)；如果 \(sum>x\)，则持续弹出堆顶元素直至 \(sum\leq x\)，再设置 \(pre_i=size\)。

有人把这种算法称为反悔堆：弹出堆顶元素其实就是“反悔”之前做出的选择，贪心更新最优解的过程。

Code Implementation

#include <iostream>
#include <queue>

using namespace std;

const int MAX_N = 3e5 + 10;

int a[MAX_N];
int pre[MAX_N], suf[MAX_N];

priority_queue<int> heap;

bool check(long long ans, int n, int k) {
    while (!heap.empty())   heap.pop();
    long long sum = 0;
    for (int i = 1; i <= n; ++i) {
        heap.push(a[i]);
        sum += a[i];
        while (sum > ans) {
            sum -= heap.top();
            heap.pop();
        }
        pre[i] = heap.size();
    }

    while (!heap.empty())   heap.pop();
    sum = 0;
    for (int i = n; i >= 1; --i) {
        heap.push(a[i]);
        sum += a[i];
        while (sum > ans) {
            sum -= heap.top();
            heap.pop();
        }
        suf[i] = heap.size();
    }

    for (int i = 1; i < n; ++i) {
        if (pre[i] + suf[i + 1] >= k)   return true;
    }
    return false;
}

long long solve(int n, int k) {
    long long l = 0, r = 0;
    for (int i = 1; i <= n; ++i)    r += a[i];

    while (l < r) {
        long long mid = l + r >> 1;
        if (check(mid, n, k))
            r = mid;
        else
            l = mid + 1;
    }   
    return r; 
}

int main() {

    int t;
    cin >> t;

    while (t--) {
        int n, k;
        cin >> n >> k;

        for (int i = 1; i <= n; ++i)    cin >> a[i];

        cout << solve(n, k) << endl;
    }   

    return 0;
}