fork download
copy 
  1. #include <bits/stdc++.h>
  2. #define all(x) (x).begin(), (x).end()
  3. #define ll long long
  4. using namespace std;
  5.  
  6. int main() {
  7.     ios_base::sync_with_stdio(false);
  8.     cin.tie(nullptr);
  9.  
  10.     int t;
  11.     cin >> t;
  12.     while (t--) {
  13.         int n;
  14.         cin >> n;
  15.         vector<int> a(n);
  16.         ll mx = 0;
  17.         for (int i = 0; i < n; ++i) {
  18.             cin >> a[i];
  19.             if ((ll)a[i] > mx) mx = a[i];
  20.         }
  21.  
  22.         vector<int> b;
  23.         b.reserve(n);
  24.         for (int i = 0; i < n; ++i) {
  25.             if (a[i] <= n) b.push_back(a[i]);
  26.         }
  27.         sort(b.begin(), b.end());
  28.  
  29.         int m = 0;
  30.         for (int x : b) {
  31.             if (x == m) ++m;
  32.             else if (x > m) break;
  33.         }
  34.  
  35.         if (m == 0) {
  36.             cout << (ll)n * mx << '\n';
  37.             continue;
  38.         }
  39.  
  40.         ll best = 0;
  41.         for (int k = 0; k <= m; ++k) {
  42.             ll kk = k;
  43.             ll mm = m;
  44.             ll nn = n;
  45.             ll cur;
  46.             if (k == m) {
  47.                 cur = mm * mm + (nn - mm) * (mm + mx);
  48.             } else {
  49.                 ll sumMex = kk * kk;
  50.                 ll sumFromKtoM = (mm + kk) * (mm - kk + 1) / 2;
  51.                 cur = sumMex + sumFromKtoM + mm * (nn - mm - 1) + (nn - kk) * mx;
  52.             }
  53.             if (cur > best) best = cur;
  54.         }
  55.  
  56.         cout << best << '\n';
  57.     }
  58.  
  59.     return 0;
  60. }
  61.  
Success #stdin #stdout 0s 5320KB
comments (?)
stdin
copy 
5
5
0 0 0 0 0
2
0 1
5
1 1 1 1 0
6
1 1 4 5 1 4
1
1000000000
stdout
copy 
5
4
13
30
1000000000
https://ideone.com/JJMNFy
language:
C++ (gcc 8.3)
created:
12 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!