PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: nilayan17, awoholo
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Simple
PREREQUISITES:
Sorting
PROBLEM:
You have an array A, on which the following operation can be performed any number of times:
- Choose any non-empty subsequence of A, and replace each of its elements with \text{MEX}(A).
Find the maximum possible sum of A.
EXPLANATION:
After performing our operations, some elements will have changed, while some will be unchanged.
Suppose k elements, at indices i_1, i_2, \ldots, i_k, have been changed.
Then, it turns out that in an optimal solution, the best we can do is to set all these elements to be equal to k via a sequence of operations.
Proof
First, let’s see that we can in fact turn all these elements to (at least) k.
There’s a fairly simple scheme to do so.
First, operate on only index i_1, then only index i_2, then only i_3, and so on.
Note that after this sequence of operations, the values at indices i_1, i_2, \ldots, i_k will all be distinct - since each time an index is modified, it’s being set to the mex of the array (and the mex can’t exist in the array, by definition).This means once the above process is done, the values 0, 1, 2, \ldots, k-1 will surely exist in A - meaning the mex of A is (at least) k.
So, we can now just select all of \{i_1, i_2, \ldots, i_k\} as our subsequence, and all of them will be set to a value of \geq k.Now, consider the case when all these values are set to a value that’s \gt k.
The only way this can happen is if, among the non-modified elements, there exists a value \leq k.
But then we could include this value into the modified elements and increase it to at least k+1, which will strictly increase the sum.So, in an optimal solution, the k modified elements will be set to exactly k, as claimed.
Now that we know this, what’s the optimal k to choose?
It’s hard to say directly, but we can simply try each of them and take the best option.
So, let’s fix the value of k, and try to decide which k elements should be modified.
As noted above, all modified elements will be set to exactly k in the end.
So, their original values don’t matter at all: the only thing that matters is the values of the elements that are not modified.
With our aim being to maximize the sum, it’s of course optimal to choose the largest N-k elements to not modify.
So, with a fixed k, the best score we can get is k^2, plus the sum of the N-k largest elements of A.
The final answer is the maximum of this across all k from 0 to N.
To compute this quickly, note that you can just sort A in descending order first, and then define S_k to be the sum of the largest k elements of A (so that S_k = A_k + S_{k-1} after sorting).
The answer is then the maximum of k^2 + S_{N-k} across all k.
TIME COMPLEXITY:
\mathcal{O}(N\log N) per testcase.
CODE:
Author's code (C++)
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
int main() {
ll i,j,k,t,n,m,f,q;
cin >> t;
while(t--){
cin >> n;
vector<ll>a(n);
for(i=0;i<n;i++)
cin >> a[i];
sort(a.begin(),a.end());
vector<ll>suff(n+1,0);
for(i=n-1;i>=0;i--)
suff[i]=suff[i+1]+a[i];
ll ans=suff[0];
for(i=0;i<n;i++)
ans=max(ans, (i+1)*(i+1)+suff[i+1]);
cout << ans << endl;
}
return 0;
}
Editorialist's code (PyPy3)
for _ in range(int(input())):
n = int(input())
a = sorted(list(map(int, input().split())))
ans = sum(a)
suf = sum(a)
for i in range(n):
suf -= a[i]
ans = max(ans, (i+1)*(i+1) + suf)
print(ans)