【每日一题】补档 CF730I. Olympiad in Programming and Sports | 反悔贪心 | 困难

题目内容

原题链接

给定 $n$ 个学生，第 $i$ 个学生的编程能力为 $a_i$ ，运动能力为 $b_i$ 。

现在要选择 $p$ 个学生进入编程队，$s$ 个学生进入运动队，每个学生只能加入一个队。

问加入两个队的 $p+s$ 个学生的能力之和最大是多少。

注意： 加入运动队的学生能力为其运动能力，加入编程队的学生能力为其编程能力。

数据范围

• $2\le n\le 3000$
• $1\le a_i,b_i\le 3000$
• $p+s\le n$

题解

我们考虑这么一个问题：假设当前已经选择过了 $p+s$ 个学生，其中 $p$ 个加入了编程队，$s$ 个加入了运动队。

但是，现在存在一个没有加入任意一个队伍的学生，他的编程能力大于编程队中某个学生的编程能力。

这符合现状吗？显然是不符合的，所以我们可以考虑先凑齐一个队伍。

即我们将所有学生按照其编程能力从大到小排序，选择前 $p$ 个学生加入编程队。之后，我们来考虑选择 $s$ 个学生加入运动队。

我们的目的是使得所有的学生在对应队伍的能力值之和最大。

所以接下来我们选择运动队的学生，有两种选择方式：

• 直接选择当前既不在编程队，也不在运动队中的，运动能力最高的学生，加入运动队
• 从编程队中，选择一个 $b_i-a_i$ 最大的学生，将其从编程队转移到运动队，然后从既不在编程队，也不在运动队中的，编程能力最高的学生，加入编程队。

这样始终保证编程队的人数为 $p$ ，每次选择都使得运动队的人数加 1 。

每次我们选择上述两种方式中，使得总能力变得更大的一种选择。

因此需要维护三个堆：

• 在编程队中的学生，维护一个 $b_i-a_i$ 的大根堆
• 既不在编程队，也不在运动队中的学生，维护一个 $a_i$ 的大根堆
• 既不在编程队，也不在运动队中的学生，维护一个 $b_i$ 的大根堆

当然，介于这题的数据范围，完全可以两层循环 $O(n^2)$ 来解决，从而不用麻烦地维护三个堆。

时间复杂度：$O(n\log n)$

代码

#include <bits/stdc++.h>
using namespace std;

struct Node1 {
    int val;
    int idx;
    Node1(int val, int idx): val(val), idx(idx) {}
    bool operator< (const Node1& A) const {
        return val < A.val;
    }
};

struct Node2 {
    int val;
    int idx;
    Node2(int val, int idx): val(val), idx(idx) {}
    bool operator< (const Node2& A) const {
        return val < A.val;
    }
};

void solve() {
    int n, p, s;
    cin >> n >> p >> s;

    vector<int> vp(n), vs(n);
    vector<array<int, 3>> vec(n);
    for (int i = 0; i < n; ++i) { cin >> vp[i]; vec[i][0] = vp[i]; } ;
    for (int i = 0; i < n; ++i) { cin >> vs[i]; vec[i][1] = vs[i]; vec[i][2] = i; };

    sort(vec.begin(), vec.end(), [](const auto& A, const auto& B) {
        return A[0] > B[0];
    });

    vector<int> vis(n, 0);

    priority_queue<Node1> heap_p2s;
    int ans = 0;
    for (int i = 0; i < p; ++i) {
        ans += vec[i][0];
        heap_p2s.emplace(vec[i][1] - vec[i][0], vec[i][2]);
        vis[vec[i][2]] = 1;
    }

    priority_queue<Node2> last_p, last_s;
    for (int i = p; i < n; ++i) {
        last_p.emplace(vec[i][0], vec[i][2]);
        last_s.emplace(vec[i][1], vec[i][2]);
    }

    // 两种选择：
    // 要么从 p2s 队列中选一个 s - p 最大的，转移到 s 队列中，然后选一个编程能力最大的
    // 要么直接从剩下的部分选一个运动能力最大的
    // 两者选一个

    for (int c = 0; c < s; ++c) {
        while (!last_p.empty() && vis[last_p.top().idx]) last_p.pop();
        while (!last_s.empty() && vis[last_s.top().idx]) last_s.pop();

        // 方案1，选一个 p2s 中 s-p 最大的
        if (!heap_p2s.empty()) {
            auto top = heap_p2s.top();
            int v1 = top.val + last_p.top().val;
            int v2 = last_s.top().val;
            if (v1 <= v2) {
                auto top2 = last_s.top();
                vis[top2.idx] = 2;
                ans += v2;
            } else {
                auto top2 = last_p.top();
                vis[top.idx] = 2;
                vis[top2.idx] = 1;
                heap_p2s.pop();
                heap_p2s.emplace(vs[top2.idx] - top2.val, top2.idx);
                ans += v1;
            }
        } else {
            // 方案2，选一个运动能力最大的
            auto top = last_s.top();
            vis[top.idx] = 2;
            ans += top.val;
        }
    }

    cout << ans << "\n";
    for (int i = 0; i < n; ++i) {
        if (vis[i] == 1) cout << i + 1 << " ";
    }
    cout << "\n";
    for (int i = 0; i < n; ++i) {
        if (vis[i] == 2) cout << i + 1 << " ";
    }
    cout << "\n";
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int T = 1;
    //cin >> T;
    while (T--) {
        solve();
    }

    return 0;
}