PROBLEM LINK:
Author: Vitaliy Herasymiv
Tester: Shiplu Hawlader
Editorialist: Praveen Dhinwa
DIFFICULTY:
Medium
PREREQUISITES:
expectation value, dynamic programming, simple combinatorics.
PROBLEM:
Given n blocks such that i th block has width A i and color C i . You have randomly arranged the blocks in some order,
then you count the number of pairs i, j such j - i = k and Color corresponding to i = Color corresponding to j. Let us call such a pair
of i, j a “good” pair.
Find out the expected number of such “good” pairs i, j for all possible permutations of n blocks.
Constraints
Note that maximum number of blocks having same color could be at most 2. Number of blocks n will <= 20. Sum of array A will be less than 200.
There are 10 test cases also.
QUICK EXPLANATION
-
As we can not try all possible permutations because they are huge(20!). We have to think of some other strategy.
-
So we will use linearity of expectation and try to count some other quantity which could lead us to the final required answer.
-
So instead of trying all permutations, we will try all “good” pairs and for each pair, we will find out the possible number of permutations
in which they could appear. Summing over (number of “good” pairs * number of permutations) will give us total number of “good” pairs. Dividing
it by total number of permutations (n!) will give us the expected value of number of “good” pairs. -
First count number of pairs i, j (j - i = k) in the single block A x .
-
Then we will try to calculate number of good pairs for each pairs of blocks A i , A j with same
color (C i = C j ).
Let us say that they have L width between them in which we will try to fit other blocks.
We need to find out number of possible permutations for such cases. We can apply simple dynamic programming algorithm to count number of ways of
fitting count elements in the given width. -
The constraint that maximum number of blocks having same color could be at most 2, ensure that we don’t need to consider triplets of same color
because they wont occur in our problem due to constraints. -
Add the above two quantities and divide them by total number of permutations i.e. (n!) and we will get the expected number of pairs.
EXPLANATION
As n <= 20, so we can’t try all possible permutations of n and check the number of good pairs, we need to do something different.
Let us do the reverse, we will try to fix some blocks having same color
and then try to find out number of possible permutations satisfying this condition. This is also called [linearity of expectation][20].
-
First we will count number of good pairs for single blocks. For a single block if it has size < K, then there can not exist two indices i, j
such that j - i = K.
For a block of length A i , the number of pairs j - i = k will be exactly A i - K.
Proof:
Fix i, i + K <= A i . So i <= A i - K. As i goes from 1 to A i - K. So there are exactly A i - K pairs. -
Now we will try to count number of pairs for blocks A i , A i having same colors. For checking the condition j - i = K,
we can see that
only condition that matters is the width of the blocks which will be between them. So we will iterate over possible separation between
the blocks from L = 1 to 200 (Max possible value of separation).For a fixed A i , A i and fixed separation L between them, we can find out the number of possible permutations satisfying
this condition.Now in the fixed separation of length L, there will be some blocks which will fit into this gap. We only need the count of such blocks fitting in
the width L.So we will iterate over count where count denotes the number of the blocks fitting in the separation. count will go from 1 to n - 2.
-
So now we need to find number of ways of selecting k blocks from A 1 to A n (excluding A i and A j )
such that their sum of widths is equal to L.
We can solve this problem by dynamic programming algorithm.Let dp[idx][count][total_width] denote the dp state. Here idx denotes the index at which we currently are, count denotes number of selected blocks.
Let total_width denote the sum of widths of the blocks of the count selected blocks.- idx will go from 1 to n.
- count will go from 1 to n - 2.
- total_width will go from 1 to 200.
Transitions of the dp:
- from state (i, count, total_width) you can go to following states. - You can select the current block and go to state dp(i + 1, count + 1, total_width + A_i). - You can skip the current block and go to start dp(i + 1, count, total_width);
Pseudo code:
for (int idx = 0; idx < n; idx++) {
// idx is equal to i or j skip.
if (idx == i || idx == j) continue;
for (int count = 0; count <= idx + 1; count++) {
// here L denotes the width between the blocks.
for (int L = 0; L <= 200 - A[i] - A[j]; L++) {
int new_L = L + A[idx];
if (new_L <= 200)
// take the current element.
dp[idx + 1][count + 1][new_L] += dp[idx][count][L];
// skip
dp[idx + 1][count][L] += dp[idx][count][L];
}
}
}
-
Number of permutations with A i and A j being a pair with equal colors and L being separation between them is
count! * (n - 2 - count)! * (n - 1) (n - 1 denotes the number of ways of starting of the A i , A j . the first block A i can start at n - 1 possible
positions.). You can easily work out this expression yourself. -
For given two pairs of blocks A i and A j with equal colors and having L separation between them, we can easily count number of pairs of position
x, y such that y - x = K and both have same color.
For this you can simple iterate over possible x positions in block A i and check whether x + K fits in the second block A j or not.
Pseudo Code:
int cnt = 0;
for (int x = 1; x <= A[i]; x++) {
if (x + K > A[i] + L && x + K <= A[i] + L + A[j]) {
cnt ++;
}
}
Complexity:
For each test case, You can have at most 10 pairs of A i , A j such that color of A i and A j is equal.
For each such pair the dp algorithm will take
O(n * m * 200).
So overall time complexity is 10 * n * m * 200 <= 10 * n * n * 200 = 2000 * n * n = 2000 * 20 * 20 = 8* 10^6.
So for 10 test case, total number of operations will be around 10^8 which will pass under the time limit easily.
