TC SRM 629: A solution for div1 950pt

№	Пользователь	Рейтинг
1	tourist	4009
2	jiangly	3823
3	Benq	3738
4	Radewoosh	3633
5	jqdai0815	3620
6	orzdevinwang	3529
7	ecnerwala	3446
8	Um_nik	3396
9	ksun48	3390
10	gamegame	3386

№	Пользователь	Вклад
1	cry	167
2	Um_nik	163
3	maomao90	162
3	atcoder_official	162
5	adamant	159
6	-is-this-fft-	158
7	awoo	157
8	TheScrasse	154
9	Dominater069	153
9	nor	153

This was originally intended as an answer for this comment, but eventually I realized that it's too long (also, the Preview feature takes way too long to process it), so it's a blog post instead.

This has complexity of $\text{[math]}$ and a horrible constant. If there's a better solution, write it in the comments.

UPD: KADR provides a solution using polynomial interpolation in O(K²).

Problem statement

There are N boxes numbered 1..N; box number i contains i candies and 1 stone. We pick 1 random object from each box. Calculate $\text{[math]}$ , where p is the probability that K of these N objects are candies.

Constraints: 1 ≤ K ≤ N ≤ 10⁹, 2·10⁹ ≥ M ≥ 10⁹ and M is prime.

Solution

It's clear that since the probability of picking a candy from box k is $\text{[math]}$ and the probability of not picking one is $\text{[math]}$ , the answer is (S_n denotes the set $\text{[math]}$ )

$\text{[math]}$

Let's denote this sum as P(N, K) and introduce another sum (N / 2 is integer division)

$\text{[math]}$

which calculates the same sum P(N, K), but with an additional restriction that S can only be subsets of range $\text{[math]}$ .

If that restriction was instead that S can only be subsets of range $\text{[math]}$ , the sum would obviously be just P(N / 2, K).

Also, let's denote $\text{[math]}$ .

Naive solution

If the constraints were small, we could use simple DP on P(N, K). If we don't use N in S, then we sum up π(S') of the same subsets S' as in P(N - 1, K). If $\text{[math]}$ , we only need K - 1 elements from the range S_N - 1, so we sum up π(S) = Nπ(S') of the same subsets S' as in P(N - 1, K - 1).

This runs in O(NK) and would obviously give TLE. What we can do instead is splitting the range S_N into 2 halves of size N / 2 (and possibly 1 integer N / 2 + 1 in the center for odd N).

Calculating P(N, K) — even N

Suppose we knew how to calculate P₂(, ) already. In order to calculate P(N, K) for even N, we can split any $\text{[math]}$ disjointly and uniquely into $\text{[math]}$ and $\text{[math]}$ .

If |S_a| = i, then |S_b| = K - i; any 2 sets S_a and S_b can be merged to form one set S and so we have for fixed i

$\text{[math]}$

Since i can be any integer in [0, K], we get

$\text{[math]}$

Calculating P(N, K) — odd N

In this case, we can again split S disjointly and uniquely into S_a, S_b as above and another set $\text{[math]}$ . We have 2 choices for S_c: either it's empty or contains N / 2 + 1; if it's empty, then we get the same sum as for even N.

If S_c is non-empty, we can again use that π(S) = π(S_a)π(S_b)π(S_c) = π(S_a)π(S_b)(N / 2 + 1), where we can take N / 2 + 1 out of the sum and get a similar sum as for even N — the only change is that if |S₁| = i, then |S₂| = K - 1 - i and

$\text{[math]}$

Again iterating over all i from 0 to K - 1, we get a formula for odd N:

$\text{[math]}$

Calculating P₂(N, K)

Let's expand the sum as

$\text{[math]}$ $\text{[math]}$

It's clear that if |S'| = K - i, then

$\text{[math]}$ $\text{[math]}$

where $\text{[math]}$ denotes the number of sets S (of size K) satisfying the condition.

How to count that number? We can just exclude set S' from all S and S_N / 2 (since they all contain S'); then, we can see that we just need to count the number of subsets of S_N / 2\ S' of size i. That's obviously just $\text{[math]}$ , so

$\text{[math]}$ $\text{[math]}$

Binomial coefficient

The binomial coefficient has kinda big arguments, no? We can't just pre-compute a Pascal triangle or factorials, but i is sufficiently small, so we can just use one of the most basic formulas for binomial coefficients

$\text{[math]}$

and so, for given N, K, compute the necessary binomial coefficients along with corresponding terms in the sum for P₂(N, K). It's useful to have modular inverses precomputed; here, we can use that M is a prime larger than K, so by Fermat's little theorem, $\text{[math]}$ .

Complexity

We now have a fairly simple way to compute P(N, K) and P₂(N, K) in O(K) time with some pre-computation. The important thing is that thanks to integer division, the N in the argument can only be N from the input divided by powers of 2; since there are just $\text{[math]}$ such powers that don't lead to the trivial case N = 0, and with fixed N, we only spend O(K²) time on computing P() and P₂() for all possible K, the complexity is $\text{[math]}$ .

The limits are too tight, though — the biggest problem is there are a lot of modulos that take a lot of time. One possible improvement is only taking the modulo when computing the sums for P() and P₂() after every 4 additions, that gives me worst-case runtime of around 2.1 seconds (so close TLE T_T). Also, we can precompute $\text{[math]}$ , which saves us some modulo operation compared to precomputing (N + 1)ⁱ and $\text{[math]}$ separately. Code

Комментарии (22)

Написать комментарий?

Alex_2oo8

10 лет назад, # |

← Rev. 2 →

+27

There is one small optimization about modulos that can reduce the runtime a lot.

Suppose we have to compute something like this (here 0 ≤ a_i, b_i < MOD and MOD < 2³¹):

long long res = 0;
for (int i = 0; i < n; i++) {
    res += a[i] * b[i];
    res %= MOD;
}

The trick is that we can compute the sum modulo MOD² without modulo operations itself:

long long res = 0;
for (int i = 0; i < n; i++) {
    res += a[i] * b[i];
    if (res >= MOD * MOD) res -= MOD * MOD;
}
res %= MOD;

→ Ответить

Xellos

10 лет назад, # ^ |

← Rev. 3 →

Yes, that's one of many things I tried — but didn't get a noticeable increase in speed.

UPD: I played with it a bit more and what works is combining the 2 tricks — when I take A[][] and A2[][] modulo $\text{[math]}$ , I can afford to check for subtraction only on every second addition:

if(i&1 && A[s][j] >= 2*modSq) A[s][j] -=2*modSq;

and this gets AC in less than 1.95 seconds worst-case runtime.

UPD2: This is still a bit labile, since swapping arguments in the if() results in runtime close to 2 seconds and possible TLEs (when very unlucky). But it's still nothing compared to how worse the runtime gets when moving the line decP =(decP*((N>>(s+1))-j+i+1))%mod; (btw this one takes about half of the whole runtime).

Edvard

It was very hard but I'm also do it (Code 1.91s)

KADR

+77

There is another O(K²) solution that involves interpolation. First, one can notice that the answer can be calculated using the following recurrence relation:

f(N, K) = N·f(N - 1, K - 1) + f(N - 1, K)

Another observation is that f(N, K) is a polynomial of N of degree 2K. This can be proven by induction (although I didn't try to do this during the contest). So we can calculate f(i, K) for i from 1 to 3K + 1 using the above formula and then use f(K, K), f(K + 1, K), ..., f(3K + 1, K) as the values of our polynomial p(x) in points K, K + 1, ..., 3K + 1. Interpolation can be done in O(K²). Finally, the answer to the problem is p(N).

Temirulan

+14

Very good solution. But how you observe this fact?

Well, it was just guess based on the general look of the formula and on the fact that K is small.

ikbal

I did not understand why f(N, K) 's degree is 2K not K.

Let p(N, K) degree of f(N, K). Then p(N, K) = max(1 + p(N - 1, K - 1), p(N - 1, K)) right?

Could you tell what did i miss?

Consider a simpler formula:

f(0) = 0 f(N) = N + f(N - 1)

According to your logic p(N) = max(1, p(N - 1)) = 1, which is false.

So, $\text{[math]}$ and f(0, K) = 0 (K > 0). Suppose that $\text{[math]}$ , which holds for K - 1 = 0; then

where the inner sum is something like $\text{[math]}$ , which is known to be a polynomial of degree k + 2. Picking k = 2K - 2, we get that f(N, K) is a polynomial of degree 2K.

caioaao

Which algorithm you use for polynomial interpolation? Do you have a reference for implementing it in C++?

Shouldn't straightforward Lagrange be sufficient? It's basically "copy formula from Wikipedia if you don't remember it".

And it works in O(degree²).

Thanks!

YuukaKazami

+24

It is surprising for me that this problem is exactly the same as a problem which I set up 2 years ago.

problem link ( In Chinese

That time my proposed solution is O(k^2) with fascinating application of Principle of Inclusion and Exclusion. But some people came up with O(k^2logn) approach (which can be speed up to klogklogn by FFT)...

I have write a solution for that problem that time and give the insight of the this problem, Unfortunately it is in Chinese :( link

The $\text{[math]}$ solution seems to be similar to mine (based on what I could guess from the formulas). What's the idea behind your O(K²) solution?

Use Inclusion and Exclusion we can turn this problem in to counting the number of n vertices connected graph with odd or even edges...

Actually, those constraints are (i and j should not be the same), so if we flip it then we should let (i and j is the same), we can do it by connect i and j with an edge, then a connected component of a graph should be the same, and the number of the edge is the number of flip we have made, base on whether it is even or odd, we add or subtract it.

If k vertices are the same, what we want is \sum_{i=1}^{n} i^k , which can be calcualted by some formula.

darkshadows

I saw this formula for Sterling numbers of first kind. http://upload.wikimedia.org/math/e/4/7/e47bc8bab37c7a86cc15eb5ecef53f5c.png This formula works in O(K*K), but it might give 0 when take modulo with MOD due to (2n-m) factorial.

But it seems to me that it works in O(N²) what we need to compute ~~is s(N, K).~~

Won't S(n+1,n+1-k) give the answer, which can be computed in O(K*K).

Oh, right. The recurrence is a bit different.

It's strange that this formula gives a polynomial in K, not in N... but it should work.

A more relevant question about the factorials is: how do we compute huge ones? If 2n - m is just a bit smaller than the prime modulo, we can't compute it iteratively because of time (and memory) limits.

← Rev. 5 →

For S(n, m) we need $\text{[math]}$ which in case of S(n + 1, n + 1 - k) will be $\text{[math]}$ which can be done in O(k).

All other factorials similarly won't exceed O(k), IMO.

Well, that means we don't need to care about the factorials being divisible by the modulus, since we don't need modular inverses there — just multiplication.

Yes, but for n + k + 1 exceeding MOD, it'll always evaluate to 0 and that was the initial problem that I had faced in solving this problem.

Блог пользователя Xellos

Problem statement

Solution