小鸾的战争（算法设计与分析）

描述

小鸾是一个杏花帝国的君主。

有一天，杏花帝国的邻国铃音帝国向杏花帝国宣战了！

小鸾知道，铃音帝国的君主灵儿是一个非常具有战略意识的女孩子，所以她必须要抓紧时间快速组建一支听她指挥、能打胜仗的军队。

小鸾很快得到了需要士卒的作战申请。申请作战的士卒们来自许多不同的门派。

小鸾一共得到了 nn 份作战申请，也就是说有 nn 个士卒申请加入战斗。第 ii 个士卒来自编号为 cici​ 的门派，拥有 wiwi​ 的战斗力。

小鸾一共需要从这 kk 个士卒中挑选出 kk 位参与作战。

小鸾发现，因为需要应对战场上不同力量的攻击，当这些士卒们来自越多的门派时，更处理不同攻击的能力越强。精通分析计算的小鸾很快得到了一个计算军队总体实力的公式：

S=diff2+∑i=1kwiS=diff2+∑i=1k​wi​

其中，SS 为军队最终的总体实力，diffdiff 为军队中的士卒来自的不同门派数量。

小鸾想要知道，她的军队最多能拥有多少总体实力呢？

由于这个年代没有计算机，小鸾穿越时空把这个问题交给了你，希望你能尽快为她求出答案。

（为了避免时空管理局的审查，你的程序必须要在 1s1s 以内计算出小鸾问题的答案哦）

出题者：韩可可同学

输入

输入的第一行是两个用空格隔开的整数 nn 和 kk，分别为申请作战的士卒数量和小鸾需要从中挑选出的数量。

第 22 至第 n+1n+1 行，每行两个整数 cici​ 和 wiwi​，分别表示第 ii 位士卒来自的门派和其战斗力。

输出

输出一行一个整数，表示军队总体实力的最大值。

输入样例 1 

5 3
1 1
2 1
1 2
3 1
4 1

输出样例 1

13

提示

1≤n≤105,1≤k,ci≤n,0≤di≤1091≤n≤105,1≤k,ci​≤n,0≤di​≤109。

AC代码

#include <iostream>
#include <vector>
#include <algorithm>
#include <map>
#include <set>
#include <queue>

using namespace std;

struct Soldier {
    int c, w;
};

bool cmp(const Soldier &a, const Soldier &b) {
    return a.w > b.w;
}

int main() {
    int n, k;
    cin >> n >> k;
    vector<Soldier> soldiers(n);
    for (int i = 0; i < n; ++i) {
        cin >> soldiers[i].c >> soldiers[i].w;
    }

    sort(soldiers.begin(), soldiers.end(), cmp);

    long long sum = 0;
    long long ans = 0;
    int diff = 0;
    map<int, int> freq;
    priority_queue<int, vector<int>, greater<int>> dupPower; // 可替换的重复门派战斗力最小堆
    set<int> usedCategory;

    vector<Soldier> rest;

    for (int i = 0; i < k; ++i) {
        sum += soldiers[i].w;
        freq[soldiers[i].c]++;
        if (freq[soldiers[i].c] == 1) {
            diff++;
        } else {
            dupPower.push(soldiers[i].w); // 重复门派可替换
        }
    }

    for (int i = k; i < n; ++i) {
        rest.push_back(soldiers[i]);
    }

    sort(rest.begin(), rest.end(), cmp);

    ans = 1LL * diff * diff + sum;

    for (auto &s : rest) {
        if (freq[s.c] == 0 && !dupPower.empty()) {
            int removed = dupPower.top();
            dupPower.pop();
            sum -= removed;
            sum += s.w;
            diff++;
            freq[s.c]++;
            ans = max(ans, 1LL * diff * diff + sum);
        }
    }

    cout << ans << endl;
    return 0;
}