Dynamic Programming, Combinatorics
You are given balls of N different colours. There are C[i] balls of colour i. Find how many ways are there to arrange the balls in a row such that no two consecutive balls are of the same colour.
Let us consider inserting balls into the row color by color. For example, if we have C = 5, C = 3, etc. Our insertion scheme will have patterns looking like:
11111, followed by
21112211 followed by etc.
Now, we consider the space between two consecutive balls. It is either “bad” if they are of the same color, or “good” if they are of different colors. Let us denote “good” spaces with ‘+’, and “bad” spaces with ‘-’. Again taking the previous example, our scheme of insertion looks like:
Initially: + (empty “good” space)
Thus, we finally want to ask the question, “After inserting all N colors of balls into the row, we want there to be 0 bad spaces. How many ways are there of doing this?”
We can answer this question using dynamic programming. Indeed, let DP[x][y] := The number of ways of inserting balls of the first ‘x’ colors in such a way that there are exactly ‘y’ bad spaces.
We have that initially, DP = 1, DP[y] = 0 for y > 0.
Then, given that we have filled in colors upto x, (i.e. we have filled in DP[x][y] values for all y), we try to insert C[x+1] balls. Note that the total number of spaces prior to insertion is 1 + C + C + … + C[x] (lets denote this sum by spaces[x]). From a DP[x][y] state, we have y “bad” spaces and spaces[x]-y (lets denote this by ‘z’) “good” spaces.
We choose i good spaces and j bad spaces into which to insert C[x+1] balls (naturally, C[x+1] >= i+j). What will this state look like? Each ball that is inserted into a “bad” space will destroy that bad space. Each ball that is inserted into a “good” space will remain good. Each of the other balls will create a “bad” space upon insertion. Hence, the number of bad spaces will become (y - j) + (C[x+1] - i - j).
We finally need to measure how much will DP[x][y] contribute to the new state. We first choose i good spaces from z (=spaces[x]-y), and j bad spaces from y : This brings about a contribution of Comb(z, i) * Comb(y, j). (Here, Comb(n, r) is the number of ways of choosing r objects from a set of n objects: = n!/r!(n-r)!).
Just choosing the possible locations is not enough however. In our example, we need to have from “+1-1-1-1-1+” case, two different cases “+2-2+1-1-1+2+1-1+” and “+2+1-1-1+2-2+1-1+”. Just by choosing which good and bad spaces to insert balls would treat the above two as the same.
The reason is, we need to ensure that we are inserting C[x+1] balls into these (i+j) spaces. The question we are left to ask is: “How many ways are there of inserting C[x+1] balls into (i+j) spaces?” This can be answered by finding the number of solutions to a1 + a2 + … + ai+j = C[x+1], ai >= 1. This is simply Comb(C[x+1] - 1, i+j-1). Thus, we get an overall contribution of Comb(z, i) * Comb(y, j) * Comb(C[x+1]-1, i+j-1).
Pseudocode for the update follows:
// Filling in values of DP[x][y] from values DP[x-1][y] for(int ij = 0; ij <= C[x]; ij++) //ij = i+j as described for(int y = 0; y <= spaces[x-1]; y++) if(DP[x-1][y] > 0) for(int j = 0; j <= min(ij, y); j++) i = ij-j; z = spaces[x-1] - b; DP[x][(y-j) + (C[x]-ij)] += DP[x-1][y] * Comb[y][j] * Comb[z][i] * Comb[C[x]-1][ij-1]; // (Remember to calculate Modulo 1E9 + 7)
Finally, output DP[N];
The time complexity of the above code is O(C[x] * sum * C[x]), where sum = C + C + … C[N] >= spaces[x]. Summing this over all x gives O(sum * (C^2 + C^2 + …)) = O(sum^3).
Can be found here
Can be found here