PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: Akshat Anil Jain
Tester: Tejas Pandey
Editorialist: Nishank Suresh
DIFFICULTY:
1912
PREREQUISITES:
None
PROBLEM:
Given an array A containing N distinct integers ranging from 1 to 2N, perform the following operation exactly K times:
- Choose some 1 \leq x \leq 2N such that x \notin A, append x to A, and add \max(A) - x to your score
What is the maximum possible final score?
EXPLANATION:
Let’s look at how the score changes when we add a new element x to A. There are two cases:
- If x \lt \max(A), then we get \max(A) - x added to the score
- Otherwise, we must have x = \max(A), in which case we add 0 to the score.
Ideally, we’d like as many operations of the first type as possible, since they’re profitable. Note that the score of the first operation is maximized when x is as small as possible and \max(A) is as large as possible.
This gives us a strategy. Let M denote the initial maximum of A, before any operations are done.
- Suppose we add K elements that do not change the maximum, i.e, they are all \lt M. Then, of course we choose the smallest K elements to insert.
- Suppose we add some element that does change the maximum.
- It’s optimal to add this element first so that all later operations have a higher score.
- It’s optimal to add as large a maximum as possible, so let’s just choose 2N to add in the first move
- The other K-1 elements then should be chosen to be as small as possible to maximize the score, so choose the K-1 smallest remaining elements
The answer is the maximum of both cases. Each one’s score can be computed in \mathcal{O}(N), giving us a linear time solution.
Note that we need to be able to quickly find the sum of the K smallest elements that are not in A. This can be done by maintaining a mark array that denotes which elements are already in A, then iterating across it from 1 to 2N and adding i to the sum if mark[i] == 0.
TIME COMPLEXITY
\mathcal{O}(N) per test case.
CODE:
Tester's code (C++)
#include <bits/stdc++.h>
using namespace std;
int main() {
int t;
cin >> t;
while(t--) {
int n, k;
cin >> n >> k;
long long int ans1 = 0, ans2 = 0;
int a[n];
for(int i = 0; i < n; i++) cin >> a[i];
sort(a, a + n);
vector<int> missing;
int now = 0;
for(int i = 1; i <= 2*n; i++) {
if(a[now] == i) now++;
else missing.push_back(i);
}
for(int i = 0 ; i < k ; i++) ans1 += max(0, a[n - 1] - missing[i]);
for(int i = 0 ; i < k - 1 ; i++) ans2 += max(0, 2*n - missing[i]);
cout << max(ans1, ans2) << "\n";
}
return 0;
}
Editorialist's code (Python)
for _ in range(int(input())):
n, k = map(int, input().split())
a = list(map(int, input().split()))
mark = [0]*(2*n + 1)
b = []
for x in a:
mark[x] = 1
for i in range(1, 2*n+1):
if mark[i] == 0:
b.append(i)
mx = max(a)
ans = int(-10 ** 13)
ans = max(ans, (k-1)*(2*n) - sum(b[0:k-1]))
ans = max(ans, k*mx - sum(b[0:k]))
print(ans)
