J - Maximize Diversity of an Array

You are given an array A. The diversity of the array A is defined as number of pairs i,j,i<j such that Ai≠Aj.

You want to maximize the diversity of the array. For that, you are allowed to make at most K operations on it, in any of which, you can select a particular element and change its value to any integer in the range 1 to 109, both inclusive.

Find out the maximum diversity of the array that you can obtain.

###Input:

First line will contain T, number of testcases. Then the testcases follow.
The first line of each testcase contains two integers N,K. where N denotes the length of array A.
The next line of each testcase contains N space separated integers, the i-th of which denotes Ai.
###Output:
For each testcase, output in a single line, the answer corresponding to the testcase, which should be an integer denoting the maximum possible diversity.

###Constraints

1≤T≤20
0≤K≤109
2≤N≤105
1≤Ai≤109
###Sample Input:
3
3 10
1 2 3
4 2
1 1 2 2
6 2
2 3 3 2 4 4

###Sample Output:
3
6
14

要优先改出现数目多的数目，可以用优先队列来做，注意答案爆int

#include <bits/stdc++.h>
using namespace std;
typedef long long int LL;
const int N = 1e5 + 10;
unordered_map<int, int> vis;
priority_queue<int> q;
int s[N], num[N], a[N];
int k, n, i, j, r, T, t;
LL ans;

int main() {
    ios::sync_with_stdio(false);
    cin >> T;
    while (T--) {
        memset(s, 0, sizeof(s));
        memset(num, 0, sizeof(num));
        vis.clear();
        cin >> n >> k;
        j = 1;
        ans = (LL)(n - 1) * n / 2;
        for (i = 0; i < n; i++) {
            cin >> a[i];
            if (vis[a[i]] == 0) vis[a[i]] = j++;
            s[vis[a[i]]]++;
        }
        if (k < n) {
            for (i = 1; i < j; i++) {
                if (s[i] > 1) {
                    q.push(s[i]);
                }
            }
            while (k--) {
                if (!q.empty()) {
                    t = q.top();
                    q.pop();
                    if (t > 2) q.push(t - 1);
                }
            }
            while (!q.empty()) {
                t = q.top();
                q.pop();
                ans -= (LL)t * (t - 1) / 2;
            }
        }
        cout << ans << '\n';
    }
}