fork download 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3. const long long MOD = 1000000007;
  4. int main() {
  5.     long long t;
  6.     cin >> t;
  7.     while (t--) {
  8.         long long n, k;
  9.         cin >> n >> k;
  10.         vector<long long> a(n);
  11.         map<long long, long long> check;
  12.  
  13.         for (long long i = 0; i < n; i++) {
  14.             cin >> a[i];
  15.             check[a[i]]++;
  16.         }
  17.  
  18.         // If k >= n, we can make all elements different
  19.         if (k >= n) {
  20.             cout << (n * (n - 1)) / 2 << "\n";
  21.             continue;
  22.         }
  23.  
  24.         // Calculate initial pairs that are different
  25.         long long res = (n * (n - 1)) / 2;
  26.  
  27.         // Subtract pairs that are same
  28.         for (auto it : check) {
  29.             if (it.second > 1) {
  30.                 res -= (it.second * (it.second - 1)) / 2;
  31.             }
  32.         }
  33.  
  34.         // Create multiset of frequencies > 1
  35.         multiset<long long> ans;
  36.         for (auto it : check) {
  37.             if (it.second > 1) {
  38.                 ans.insert(it.second);
  39.             }
  40.         }
  41.  
  42.         // Process k operations
  43.         while (!ans.empty() && k > 0) {
  44.             long long temp = *prev(ans.end());
  45.             ans.erase(temp);
  46.  
  47.             if (temp > 1) {
  48.                 k--;
  49.                 temp--;
  50.                 res += temp;  // Add new pairs formed
  51.             }
  52.             if (temp > 1) {
  53.                 ans.insert(temp);  // Insert reduced frequency if still > 1
  54.         	}
  55.         }
  56.  
  57.         cout << res << "\n";
  58.     }
  59.     return 0;
  60. }
Success #stdin #stdout 0.01s 5276KB
comments (?)
stdin 
3
3 10
1 2 3
4 2
1 1 2 2
6 2
2 3 3 2 4 4
stdout 
3
5
13
https://ideone.com/kDnHyj
language:
C++ (gcc 8.3)
created:
11 hours ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!