fork download
  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. ll mm = m;
  36. ll ans = mm * mm + (ll)(n - m) * (mm + mx);
  37.  
  38. cout << ans << '\n';
  39. }
  40.  
  41. return 0;
  42. }
  43.  
Success #stdin #stdout 0s 5320KB
stdin
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
5
4
13
30
1000000000