PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
TBD
PREREQUISITES:
None
PROBLEM:
You have an array of length N.
Exactly K elements must be deleted from it.
Find all possible values that the median of the resulting array can take.
EXPLANATION:
Let M = N-K denote the length of the resulting array, after deletions.
The key to solving this problem is in relating the frequencies of elements to M.
Let’s handle the case of M being odd first.
Say, M = 2x+1.
In this case, the median will be the element at position x+1 after sorting (we’re using 1-indexing.)
Let’s try to check if some element A_i can be made the median using some deletions.
Observe that:
- There must be at most x elements strictly smaller than A_i remaining after deletion.
This is because if there were more, than the (x+1)-th smallest element would be strictly smaller than A_i too, so the median cannot be A_i. - Similarly, there must be at most x elements strictly larger than A_i as well.
The reasoning is the exact same. - If both the above conditions are satisfied, the median is guaranteed to be A_i.
So, we need to end up with an array that has at most x elements both smaller and larger than A_i, after deletion.
Suppose there are S elements smaller than A_i in the original array, and L elements larger than A_i.
Then,
- If S \gt x, we need to delete (S-x) of the smaller elements in order to make A_i the median.
Otherwise, we’re not forced to delete any of them. - Similarly, if L \gt x, we need to delete (L-x) of the larger elements, otherwise there’s no forcing.
Thus, the number of forced deletions is \max(0, S-x) + \max(0, L-x).
If this value exceeds K (the allowed number of deletions), then it’s impossible to make A_i the median of the final array.
Otherwise, it’s always possible, because we’re able to reach a state where there are at most x elements both smaller than and larger than A_i - after which we can just keep deleting smaller/larger elements till K deletions are reached.
Thus, the check for A_i being a median is simple: find S and L (the number of smaller/larger elements compared to A_i), and check if
This can be done for a single A_i in linear time, so it’s \mathcal{O}(N^2) overall which is good enough.
Next, let’s consider the case of even M.
Say, M = 2x.
In this case, the median will be the element at position x after sorting (we’re using 1-indexing.)
It’s easy to see that the exact same reasoning from the odd case applies here, just with slightly changed numbers.
Specifically, A_i can be the median if and only if:
- There are at most x-1 remaining elements smaller than A_i; and
- There are at most x remaining elements larger than A_i.
So, once again we can compute the values S and L, and this time the check becomes
Make sure to print only distinct elements as the statement asks, and don’t accidentally print duplicates.
It is possible to improve the complexity of this approach to \mathcal{O}(N\log N) by sorting the values first; which makes finding the counts S and L much easier.
However, this was not needed to receive AC due to the low constraints.
TIME COMPLEXITY:
\mathcal{O}(N^2) or \mathcal{O}(N\log N) per testcase.
CODE:
Editorialist's code (PyPy3)
for _ in range(int(input())):
n, k = map(int, input().split())
a = list(map(int, input().split()))
m = n - k
medians = set()
for i in range(n):
small, large = 0, 0
for j in range(n):
if a[j] < a[i]: small += 1
if a[j] > a[i]: large += 1
reqd = 0
if m%2 == 0:
reqd = max(0, small - (m//2-1)) + max(0, large - m//2)
else:
reqd = max(0, small - m//2) + max(0, large - m//2)
if reqd <= k: medians.add(a[i])
print(*sorted(list(medians)))