Partition Count
Problem Setter: Adhyyan Sekhsaria
Problem Tester: Soham Mukherjee
Editorialist: Taranpreet Singh
Problem:
Given an array of size N and an integer K, Find out number of ways of partitioning the whole array into partitions such that no partition contains more than K unique elements.
Solution:
Since partitions are contiguous, it follows that if a partition [l, r] contains K+1 distinct elements, Then any contiguous partition which include this partition contains more than K distinct elements. This gives shape to the famous Sliding Window Algorithm But the use of which, we will determine the rightmost index r, such that all partition [l, r-1] is valid while partition [l,r] is not. we make an array last[], where last[i] store the largest index such that partition [last[i],i] is invalid while partition [last[i]+1, i] is valid.
Now, Let us understand the reasoning behind calculating this index.
Consider the array 1 2 3 2 4 and K = 2
I am simulating the process of sliding window algo.
s is a multiset. set l = N-1, r = N;
while r< N:
set.add(A[r])
while size(set) > K:remove(A[++l]);
last[r++] = l; //
Now, this becomes a Dynamic programming problem. Let dp[i] denote Number of ways of partitioning elements upto ith element inclusive. (1-based indexing)
Base Case: dp[0] = 1, dp1 = 1.
Recurrence: dp[i] = summation(dp[k], last[i]<=k < i). Suppose we want to know number of ways i elements can be partitioned, knowing the same for all previous indices. We can include ith element in all partitions from last[i] to i-1, thus, resulting in addition.
The above recurrence still gives O(N^2) solution, because of internal loop, but we can speed it up by use of prefix arrays, Thus solving problem in O(N).
Use prefix[0] = dp[0], prefix1 = dp[0]+dp1.
Now, During recurrence, instead of running loop from last[i] to i-1, Directly set dp[i] = prefix[i-1] - prefix[last[i]-1].
Also, take care to perform operations mod 1e9+7 and also handle that negative values in mod carefully.