NTT, Multi-point Evaluation
We have to find the answer modulo 786433 = 3*2^18+1, that’s the first hint that we have to use NTT
O(N2log2N) SolutionLets find the probability that **B1** is included in the expected sum, the required answer is, **B1** * **P1** (probability to select **K-1** barmen from the rest **N-1**).
The probability to select K-1 barmen from the rest N-1, can be calculated using NTT, we simply have to calculate the power of K-1 in the polynomial,
(1 - P2 + P2x) * (1 - P3 + P3x) * … * (1 - Pn + Pnx)
This can be calculated in O(Nlog2N), using Divide and Conquer with NTT technique.
Finally, we calculate the probability that we include Bi for all i , and sum them.
O(Nlog2N + MOD*logN) Solution
Let us define F(x) = (1 - P1 + P1x) * (1 - P2 + P2x) * (1 - P3 + P3x) * … * (1 - Pn + Pnx).
The required answer is thus sum of Bi * Pi * (coefficient of Xk-1 in F(x) / (1 - Pi + Pix))
= sum Bi * Pi / (1 - Pi) * (coefficient of Xk-1 in F(x) / (1 + Hix)), where Hi = Pi / (1 - Pi)
= sum Bi * Pi / (1 - Pi) * (coefficient of Xk-1 in F(x) * (1 - Hix+Hi2x - Hi3x …))
= sum Bi * Pi / (1 - Pi) * (coefficient of Xk-1 in (1 + C1x+C2x2 + C3x3 …) * (1 - Hix+Hi2x - Hi3x …)), where Ci is coefficient of xi in F(x).
= sum Bi * Pi/ (1 - Pi) * (C0 * (-1)K * HiK + C1 * (-1)K-1 * HiK-1 …)
= sum Bi * Pi / (1 - Pi) * G(Hi), where G(x)= (C0 * (-1)K * xK + C1 * (-1)K-1 * xK-1 …)
= B1 * P1/ (1 - P1) * G(H1) + B2 * P2 / (1 - P2) * G(H2) + …
The polynomial F(x) can be constructed in Nlog2(N) using Divide and Conquer + NTT and thus G(x) can be constructed in O(N).
G(Hi) can be calculated for all i in MOD*log(N) using NTT.
Author’s Solution can be found here