PROBLEM LINKS
DIFFICULTY
MEDIUM
EXPLANATION
Consider a particular coloring. Let’s call a pair of links good if they are of the same color and there are less than M other links between them, so that there exists a block of M+1 links containing this pair. It’s easy to see that the coloring satisfies the restrictions if every block of M+1 consecutive links contains exactly one good pair. Even more, suppose we have fixed a set of good pairs and there is exactly one good pair in each block of size M+1. Then there are M! satisfying colorings having exactly such set of pairs. We can see it when we try to color the links from left to right – the coloring becomes uniquely determined after we color the first M+1 links.
The problem is thus reduced to finding the number of satisfying sets of good pairs and multiplying this number by M!. Let’s consider two arrays A and B of size K, where K is the number of good pairs. Ai and Bi are the indices of the chains of good pair i so that Ai < Bi. Let’s also sort the pairs in increasing order of Bi. For any satisfying coloring the following conditions must be held:
- B1 ≤ M+1 and (if K > 1) B2 > M+1, as the first block should contain exactly one good pair;
- N-M ≤ AK < BK and (if K > 1) AK-1 < N-M, as the last block should contain exactly one good pair;
- Ai+M+1 = Bi+1 for each 1 ≤ i < K – block Ai…Ai+M contains one good pair (Ai Bi), and block Ai+1…Ai+M+1 should contain one good pair, for which the only candidate is (Ai+1 Bi+1); also, as Ai < Bi, this implies Bi+1-Bi ≤ M.
Now, let’s look at these conditions carefully. Suppose we’ve fixed array B so that all the conditions above are satisfied. What is the number of ways to fill A? Notice that A1…K-1 are determined uniquely using the third condition, and there are BK-(N-M) ways to fill AK by the second condition.
Let f[i] be the number of ways to fill array B so that the last value BK = i (we don’t consider the second condition now). By the first condition we get f 2 = f 3 = … = f[M+1] = 1. Then, by the implication from the third condition we get f[i] = sum(f[i-M]…f[i-1]) for i > M+1. The answer now is sum(f[i]*(i-(N-M))) for i > N-M (multiplication comes from the number of ways to fill AK). This description may be difficult to fully understand, so you’re encouraged to try to figure out the exact details by yourself given the main idea.
The last thing to note is that we should calculate the values of f[] in O(N) time, as O(N2) is too much to pass the time limit. One of the ways is to maintain an auxiliary array s[] so that s1 = 0 and s[i] = s[i-1]+f[i] for i > 1 (so s[i] is the sum of f1…f[i]). Then f[i] = s[i-1]-s[i-M-1], as it can be easily seen. Another way is to notice that f[M+2] = M and f[i] = 2*f[i-1]-f[i-M-1] for i > M+2. Of course, all calculations should be done modulo 109+7
SETTER’S SOLUTION
Can be found here.
TESTER’S SOLUTION
Can be found here.