AMBALLS - Editorial

Problem Link:






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[1] = 5, C[2] = 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)

Then: +1-1-1-1-1+

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[0][0] = 1, DP[0][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[1] + C[2] + … + 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][0];

The time complexity of the above code is O(C[x] * sum * C[x]), where sum = C[1] + C[2] + … C[N] >= spaces[x]. Summing this over all x gives O(sum * (C[1]^2 + C[2]^2 + …)) = O(sum^3).

Setter’s Solution:

Can be found here

Tester’s Solution:

Can be found here


How many complete graphs exist with nodes = total number of balls, and chromatic number = N? I assume this is an equivalent statement. How to approach this way? (without DP)

1 Like

why -1 is done to this expression Comb(C[x+1] - 1, i+j**-1**) ??

1 Like

@jimmykmh : at first even i didn’t get the formula but after thinking about it for some time i got it.

Here is the reason why Comb(C[x+1] - 1, i+j-1) is correct one.

consider following analogy for better understanding.
say, given n ball,you need to add these balls in K basket(here basket is actually space)now in how many ways you can do it ??

now consider a string say bbbbbDD // here “D” means divider which are equal to ((no. of space or basket)-1) and “b” represent ball.

(no. of balls|“b”|) n=9 , (number of spaces or basket)k=3

now one of the permutation will be bDbbbDbbb.
here,using ((no. of spaces or basket)-1) dividers(“D”), i can divide “b” into 3 sets

i,e number of dividers(“D”) required = k-1

1st bag contains : b

2nd bag contains : bbb

3rd bag contains : bbb

Now we know that number of distinct permutation of bDbbbDbbb is :-

9!/(2! 7!)
which is equivalent to C(9,2),now making it generalized by using n and k.

n=9 , k =3 in above explanation.
i.e C(n+k-1,k-1);

But in the given question we do not want any spaces(i+j) to remain empty i.e each space should incorporate at least one ball.This can be achieved by adding k ball to k basket so that each basket has at least one ball.Hence we are left with n-k ball to be added in k basket.

C(n+k-1,k-1) , replacing n = n-k in this formula we get.


and hence this is how formula is working.


Here is another way to get where Comb(C[x+1] - 1, i+j-1) comes from. We need to clarify the problem.

  1. We choose i-good spaces and j-bad spaces
  2. We put all C[x+1] balls into the spaces which we chose in the first step. (We put at least a ball to the spaces, that means every spaces we have chosen will have at least a ball.)

We can reduce the problem into the problem of counting positive solutions of the equation a1+a2+a3+…aij = C[x+1]. As you know we can count the number of solutions with Combination-with-repetition(nHr). In this problem, n=(i+j), r=C[x+1]-(i+j), so the number should be nHr=(i+j)H(C[x+1]-i+j). Since nHr=(n+r-1)Cr and nCr=nC(n-r), we can rewrite the relation like this.

  • (i+j)H(c[x+1]-(i+j))=(i+j+c[x+1]-(i+j)-1)Combi(c[x+1]-(i+j))=(c[x+1]-1)Combi(c[x+1]-(i+j)=(c[x+1]-1)Combi(i+j-1)

(Here Combi stands for combination. I just used the word “Combi” to distinguish combination with the array C[i])


great solution…

1 Like

thank u…

1 Like

sorry for my bad english.
This is question nice!

Links to problem (Practice and Contest) broken. Fix!

“Complete Graphs” have a chromatic number = size of graph. What exactly did you mean?

I meant connected graphs. :frowning:

For the test case N = 3, 3 1 2.
Suppose there are 6 nodes (total number of balls) and we have to connect all the nodes such that the chromatic number of the resulting graph is 3. How many such graphs are possible?

Suppose the three colors were R(3), B(1), G(2). So one possible graph is G-R-B-R-G-R

Connected graphs is still more general. Here we want paths. Also “chromatic number” of paths is 2 :slight_smile:
“Equivalent statement” is stronger than what you suggest. The implication needs to be both ways.

Fixed, thanks!

link to setter’s solution broken.Please fix it

1 Like

I don’t understand as well. Why isn’t the ways of inserting C[x+1] balls into (i+j) spaces Comb(C[x+1], i+j) ?

I’m trying to understand by tracing the pseudocode. For the first run, I get

ij = 0,

y = 0,

j = 0,

i = 0,

z = 1,

DP[1][5] += DP[0][0] * Comb[0][0] * Comb[1][0] * Comb[4][-1]

What is Comb[4][-1] or am I missing or misunderstanding something?

Thanks. I finally understood… in part thanks to you saying “we do not want any spaces(i+j) to remain empty”

How I eventually reached my understanding was for n balls and k spaces:

  1. there are n-1 places where I can insert a partition such that there will be no empty space

  2. there are k-1 partitions I need to insert such that I have k spaces

Hence C(n-1, k-1)

1 Like

Sorry for not adding it to the Editorial.

There is another way to get Comb(n-1, k-1).

First write the numbers 1, 2, 3, …, n. Then From 1…n-1, pick k-1 numbers. Lets say you’ve picked a1, a2, a3, …, a_{k-1}.

You get your original “number of balls in each bin, with no empty bin” by taking consecutive differences. So thats a1-0, a2-a1, a3-a2, …, a_{k-1}-a_{k-2}. Thats your k-1 bins. Your last bin will have n-a_{k-1} balls.