寒训5D Least Prefix Sum

题意

给定长度为n的数组a[1],a[2],...,a[n],你可以对这个数组的某一项进行取反操作。记s[i]为数组a[n]前i项的和，若想要让s[m]为数组s[n]中最小值，求取反操作次数的最小值。

思路

若s[m]为最小前缀和，则当i<m或i>m时，总有s[i] >= s[m] ，移项得到 s[i] - s[m] >= 0。

i<m时，s[i] - s[m] = -(a[i+1] + a[i+2] + ... + a[m]) >= 0

i>m时，s[i] - s[m] = a[m+1] + a[m+2] + ... + a[i] >= 0

对于违反上述情况的情形，就需要用到取反操作，考虑每个数取反的贡献，然后对贡献最大的数进行操作，一直操作直到变为符合上述情况。为了能尽快地找到贡献最大的数，考虑用堆维护。

代码

#include <iostream>
#include <queue>
using namespace std;
int T;
int n, m, a[200005];
priority_queue<int> q1;
priority_queue<int, vector<int>, greater<int> > q2;

int solve() {
    int ans = 0;
    long long sum = 0;
    for (int i = m; i > 1; --i) {
        sum -= a[i];        //这个时候我们得到了前i-1项的前缀和 - 前m项的前缀和
        q1.push(a[i]);
        while (sum < 0) {
            int p = q1.top();
            sum += p * 2;
            q1.pop();
            ans ++;
        }
    }
    while (!q1.empty()) {
        q1.pop();
    }
    sum = 0;
    for (int i = m + 1; i <= n; ++i) {
        sum += a[i];
        q2.push(a[i]);
        while (sum < 0) {
            int p = q2.top();
            sum -= p * 2;
            q2.pop();
            ans ++;
        }
    }
    while (!q2.empty()) {
        q2.pop();
    }
    return ans;
}

int main() {

    ios::sync_with_stdio(0);
    cin.tie(0);

    cin >> T;
    while (T--) {
        cin >> n >> m;
        for (int i = 1; i <= n; ++i) {
            cin >> a[i];
        }
        cout << solve() << '\n';
    }


    return 0;
}