PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Setter:Teja Vardhan Reddy
Tester:Zhong Ziqian
Editorialist:Taranpreet Singh
DIFFICULTY:
Easy-Medium
PREREQUISITES:
Combinatorics and Factorials. Modular Arithmetic.
PROBLEM:
Given an array $A$ of size $N$, Count the number of ways we can select a subset from $A$ such that the median of selected subset is also present as an element in the subset.
SUPER QUICK EXPLANATION
- Every subset of odd size has its median present in the subset, so, we can directly add $2^{N-1}$ to answer.
- For even size subset, The subset is good, if and only if middle two elements are equal.
- We can fix the left middle element, and for every possible even size of the subset, say $2*X$, try to
EXPLANATION
This problem is probably the best example of how we can use combinatorics to optimize our solution from $O(N^3)$ to $O(N)$ (except sorting).
First of all, let us formulate the definition of median.
If the size of the selected subset is odd, the median is just the middle element of subset after sorting. Since the middle element is present in the subset, all subsets of odd size are valid. It can be easily proven that there are $2^{N-1}$ such subsets. So, we can directly count these and move toward even size subsets.
If the size of the selected subset is even, the median is defined as the average of two middlemost elements after sorting. Now, say we have two middle elements $x$ and $y$, with condition $x \leq y$. Let $z = (x+y)/2$ be the median of sequence. If we write $y = x+d, d \geq 0$, we can see, $z = x+d/2$ and also, $z = y-d/2$. This way, we can see, that the median of a sequence can never be smaller than $x$ and greater than $y$. So, For $z$ to be present in subset, we need either $z = x$ or $z = y$. But, this would imply $d = 0$, Hence leading to the conclusion that
For an even size subset to be valid, the two middlemost elements should be equal. This forms the crux of our solution, and now, we need to count the number of even sized subsets with equal middle elements.
After all this, there are a number of approaches to solving this problem, all of which required us to sort the array $A$ first.
Let us consider a $O(N^3)$ solution first.
Iterate over every pair of equal elements $(i, j)$ such that $A_i = A_j$ and iterate over the size $2*X$ of subset from $X = 1$ to $N$. The number of ways to make the subset of size $X$ with two fixed middle elements is just the product of the number of ways we can select $X-1$ elements from $[1, i-1]$ and $X-1$ elements from $[j+1, N]$.
This solution requires to iterate over every pair $(i,j)$ which takes $O(N^2)$ time and $O(N)$ time per pair, leading to Overall time complexity $O(N^3)$ which shall pass only the first subtask.
For a faster implementation, We will see two implementations, one from setter, and another interesting solution which we can further optimize to $O(N)$ except for sorting.
Setter's Implementation
First of all, see what we were doing in the previous solution.
We were fixing two equal elements and tried to count the number of ways we can make subsets of all sizes. Now, We shall fix only the Left Middle Element (Or Right one, whichever implementation you prefer).
Suppose we fixed the ith element as the left middle element. Now, We will iterate over all sizes $2*X$ and try to include $X-1$ elements from $[1, i-1]$ and $X$ elements from $[i+1, N]$. We need the right middle element to be same as the left middle element. So, When choosing $X$ elements from $[i+1, N]$, we need at least one occurrence of $A[i]$. This is same as subtracting all the ways to select $X$ elements in the range $[i+1, N]$ which do not have $A[i]$ at all. Suppose Number of occurrence of $A[i]$ in range $[i+1, N]$ is $f$, then we can count the number of ways to select $X$ elements from range $[i+1, N]$ such that it contains at least one occurrence of $A[i]$ as
$T$ = Number of ways to select $X$ elements out of $N-i$ elements less Number of ways to select $X$ elements out of $N-i-f$ elements.
We can select $X-1$ elements from $[1, i-1]$ in suppose $U$ ways. Then, Number of ways we can have good subsets with $A[i]$ as the left middle element is $U*T$. Summation of this product for all sizes for all elements gives us the number of good subsets of even size. We can add to it, the number of good odd sized subsets and print the answer.
Alternative Implementation.
For this solution, We will not count good subsets, but subtract bad subsets from total subsets. The total number of subsets is $2^N$ out of which one is empty. Empty subset doesn't contain its median, hence a bad subset. So left with $2^N-1$ subsets.
In this solution, we will count the number of subsets which have the different value of middle elements. For this solution, we iterate over all distinct elements and try to the number of ways to select subsets of all even sizes.
In this solution too, we fix the left middle element and try to count the number of bad subsets.
Suppose we have $cnt$ elements smaller than $A[i]$ and There are total $f$ occurrences of $A[i]$ and we have to make a subset of size $2*X$, Then, we need to select $X$ elements from $cnt+f$ such that it contains atleast 1 occurrence of $A[i]$ and selecting $X$ elements from Elements strictly greater than $A[i]$, THere are $N-cnt-f$ such elements. The required number of ways is the product of both. We can find the summation of this product over all even sizes of subsets for all distinct values.
This solution also has time complexity $O(N^2)$, thus fits the time limit easily.
Hint to O(N) Solution. Note: This optimization is not exactly necessary.
We rely on the same idea as the alternative implementation. We still iterate over all distinct values, but now, we will find a closed form to count bad subsets of all sizes. But the thing is, we try to find a closed form for all sizes.
We see, the summation is like $^nC_x*^mC_x$ for all sizes $x$. Can you find the closed form? Seems like many users actually did, and submitted a Linear time solution during the contest itself.
For those who did not reduce, this is left as an exercise for them.
Looking for something helpful? See Vandermonde's identity here.
Computing Number of ways to select $X$ elements out of $N$ elements
The thing is, we need to calculate N select X efficiently. We know from the definition of Combinations that we can represent them in form of factorials. We can precompute Factorials and their inverses for the division. Then, $^NC_R = N!*inv(R!)*inv((N-R)!)$.
For details, refer this excellent post here.
Time Complexity
Time complexity is $O(N^2)$ per test case for the intended solution.
AUTHOR'S AND TESTER'S SOLUTIONS:
Setter's solution
Tester's solution
Editorialist's solution
Feel free to Share your approach, If it differs. Suggestions are always welcomed. :)
