School Clubs
Contest: Codeforces Round 700 (Div. 1)
Difficulty rating: 3500
Tags: dp, fft, math, number theory, probabilities

In Homer's school, there are $$$n$$$ students who love clubs.

Initially, there are $$$m$$$ clubs, and each of the $$$n$$$ students is in exactly one club. In other words, there are $$$a_i$$$ students in the $$$i$$$-th club for $$$1 \leq i \leq m$$$ and $$$a_1+a_2+\dots+a_m = n$$$.

The $$$n$$$ students are so unfriendly that every day one of them (chosen uniformly at random from all of the $$$n$$$ students) gets angry. The student who gets angry will do one of the following things.

- With probability $$$\frac 1 2$$$, he leaves his current club, then creates a new club himself and joins it. There is only one student (himself) in the new club he creates.
- With probability $$$\frac 1 2$$$, he does not create new clubs. In this case, he changes his club to a new one (possibly the same club he is in currently) with probability proportional to the number of students in it. Formally, suppose there are $$$k$$$ clubs and there are $$$b_i$$$ students in the $$$i$$$-th club for $$$1 \leq i \leq k$$$ (before the student gets angry). He leaves his current club, and then joins the $$$i$$$-th club with probability $$$\frac {b_i} {n}$$$.

Homer wonders the expected number of days until every student is in the same club for the first time.

We can prove that the answer can be represented as a rational number $$$\frac p q$$$ with $$$\gcd(p, q) = 1$$$. Therefore, you are asked to find the value of $$$pq^{-1} \bmod 998\,244\,353$$$. It can be shown that $$$q \bmod 998\,244\,353 \neq 0$$$ under the given constraints of the problem.

Input format:
The first line contains an integer $$$m$$$ ($$$1 \leq m \leq 1000$$$) — the number of clubs initially.

The second line contains $$$m$$$ integers $$$a_1, a_2, \dots, a_m$$$ ($$$1 \leq a_i \leq 4 \cdot 10^8$$$) with $$$1 \leq a_1+a_2+\dots+a_m \leq 4 \cdot 10^8$$$, where $$$a_i$$$ denotes the number of students in the $$$i$$$-th club initially.

Output format:
Print one integer — the expected number of days until every student is in the same club for the first time, modulo $$$998\,244\,353$$$.

Constraints:
Time limit: 4.0 s
Memory limit: 512.0 MB

Note:
In the first example, no matter which student gets angry, the two students will become in the same club with probability $$$\frac 1 4$$$. So the expected number of days until every student is in the same club should be $$$4$$$.

In the second example, we note that in the first day:

- The only student in the first club will get angry with probability $$$\frac 1 3$$$. If he gets angry, then he will create a new club and join it with probability $$$\frac 1 2$$$ (In this case, there will be three clubs which have $$$0, 1, 2$$$ students in it, respectively), leave his current club and join the second club with probability $$$\frac 1 2 \cdot \frac 2 3 = \frac 1 3$$$, or stay still with probability $$$\frac 1 2 \cdot \frac 1 3 = \frac 1 6$$$;
- Each of the two students in the second club will get angry with probability $$$\frac 1 3$$$. If one of them gets angry, then he will create a new club and join it with probability $$$\frac 1 2$$$, leave his current club and join the second club with probability $$$\frac 1 2 \cdot \frac 1 3 = \frac 1 6$$$, or stay still with probability $$$\frac 1 2 \cdot \frac 2 3 = \frac 1 3$$$.

In the fourth example, there is only one club initially. That is, every student has already been in the same club. So the answer is $$$0$$$.

Example input:
2
1 1

Example output:
4

Example input:
2
1 2

Example output:
18

Editorial:
Section 1. The Expected Value of Stopping Time
There may be several ways to deal with the expected value of stopping time. Here, we decide to give an elegant way to derive the expected stopping time inspired by MiFaFaOvO's comment. Here, we are going to introduce the mathematical tool we will use in the following description (see here for a Chinese explanation.)
Consider a random process $$$A_0, A_1, A_2, \dots$$$, where $$$A_i$$$ is called the $$$i$$$-th event of the process, or the event at time $$$i$$$. Let random variable $$$T \in \mathbb{N}$$$ be its stopping time. That is, for every $$$t \in \mathbb{N}$$$, It can be decided whether $$$T = t$$$ as long as $$$A_0, A_1, \dots, A_t$$$ are given. If there is a function $$$\phi(\cdot)$$$ that maps an event to a real number such that
1. $$$\mathbb{E}[\,\phi(A_{t+1})-\phi(A_t) \mid A_0, A_1, \dots, A_t \,] = -1$$$ for every $$$t \in \mathbb{N}$$$, and
2. $$$\mathbb{E}[\phi(A_T)]$$$ depends on some properties of the events and the stopping time, which is usually a constant in practice.
To begin our journey, we should first describe an event as something (a number, a tuple, a sequence, or a set) in our problem. We find it naturally to use a (multi-)set $$$A = \{ a_1, a_2, \dots, a_m \}$$$ to describe it, where $$$m$$$ denotes the number of clubs, the $$$i$$$-th of which contains $$$a_i$$$ students. We see that the explicit order of the clubs does not matter but the number of students in each club. Furthermore, we notice that every club's status is equal, which implies us to define the potential function as the form $$$$$$ \phi(A) = \sum_{a \in A} f(a) = \sum_{i=1}^m f(a_i), $$$$$$ where $$$m = |{A}|$$$. It should be noted that the sum of elements in $$$A$$$ is a constant, so we let $$$$$$ n = \sum_{a \in A} a = \sum_{i=1}^m a_i. $$$$$$
We first consider the corner case for the value of $$$f(0)$$$. Since students will never join an empty club, $$$A$$$ and $$$A \cup \{0\}$$$ denotes the events with the same property. Therefore, we need $$$\phi(A) = \phi(A \cup \{0\})$$$, which immediately yields $$$f(0) = 0$$$.
We now investigate how the potential function behaves from the $$$t$$$-th event $$$A_t = \{ a_1, a_2, \dots, a_m \}$$$ to the $$$(t+1)$$$-th event $$$A_{t+1}$$$. We note that the student who gets angry is in the $$$i$$$-th club with probability $$$\frac {a_i} n$$$. After that, the angry student will
1. leave and create a new club with probability $$$\frac 1 2$$$. In this case, $$$A_{t+1} = \{ a_1, a_2, \dots, a_{i-1}, a_i-1, a_{i+1}, \dots, a_m, 1 \}$$$, and we have $$$$$$ \phi(A_{t+1}) = \phi(A_t)-f(a_i)+f(a_i-1)+f(1), $$$$$$
2. join the $$$j$$$-th club ($$$j \neq i$$$) with probability $$$\frac {a_j} {2n}$$$. In this case, $$$A_{t+1} = \{ a_1, a_2, \dots, a_i-1, a_j +1, \dots, a_m\}$$$, and we have $$$$$$ \phi(A_{t+1}) = \phi(A_t)-f(a_i)-f(a_j)+f(a_i-1)+f(a_j+1), $$$$$$
3. stay still in the $$$i$$$-th club with probability $$$\frac {a_i} {2n}$$$. In this case, $$$A_{t+1} = A_t$$$, and we have $$$$$$ \phi(A_{t+1}) = \phi(A_t). $$$$$$
To sum it up, we get that $$$$$$ \begin{aligned} \mathbb{E}[\, \phi(A_{t+1}) \mid A_t \,] = \sum_{i=1}^m \frac {a_i} n \Bigg[ & \frac 1 2 \big(\phi(A_t)-f(a_i)+f(a_i-1)+f(1)\big) \\ &+ \sum_{j \neq i} \frac {a_j} {2n} \big( \phi(A_t)-f(a_i)-f(a_j)+f(a_i-1)+f(a_j+1) \big) + \frac {a_i} {2n} \phi(A_t) \Bigg]. \end{aligned} $$$$$$ Let $$$\mathbb{E}[\,\phi(A_{t+1})-\phi(A_t) \mid A_0, A_1, \dots, A_t \,] = -1$$$, and then we obtain that $$$$$$ f(1) + \sum_{i=1}^m \frac {a_i} {n^2} \left[ (2n - a_i) f(a_i-1) + (n - a_i) f(a_i+1) - (3n - 2a_i) f(a_i) \right] + 2 = 0. $$$$$$ Here, we need to find a function $$$f(\cdot)$$$ that satisfies the above condition with the corner case $$$f(0) = 0$$$.
There are several possible ways to assign the value of $$$f(a)$$$. We choose the most simple one such that $$$f(1)+2 = 0$$$, which removes the constant term. Under this assignment, we have $$$f(1) = -2$$$. After this, we have to make that $$$$$$ (2n - a) f(a-1) + (n - a) f(a+1) - (3n - 2a) f(a) = 0, $$$$$$ i.e. $$$$$$ (n-a) \left[ f(a+1)-f(a) \right] = (2n-a) \left[ f(a)-f(a-1) \right]. $$$$$$ holds for every $$$a$$$.
We let $$$g(a) = f(a+1)-f(a)$$$ for $$$a \geq 0$$$. Therefore, $$$g(0) = f(1)-f(0) = -2$$$, and $$$$$$ \begin{aligned} g(a) & = \frac {2n-a} {n-a} g(a-1) \\ & = \frac {2n-a} {n-a} \cdot \frac {2n-(a-1)} {n-(a-1)} g(a-1) \\ & = \dots \\ & = \frac {2n-1} {n-1} \cdot \frac {2n-2} {n-2} \cdot \dots \cdot \frac {2n-a} {n-a} g(0). \end{aligned} $$$$$$ We find that $$$$$$ \begin{aligned} f(a) & = f(0) + g(0) + g(1) + \dots + g(a-1) \\ & = -2 \left[ 1 + \frac {P(1)} {Q(1)} + \frac {P(2)} {Q(2)} + \dots + \frac {P(1)P(2) \dots P(a-1)} {Q(1)Q(2) \dots Q(a-1)} \right], \end{aligned} $$$$$$ where $$$P(x) = 2n-x$$$ and $$$Q(x) = n-x$$$. If we can compute $$$f(a)$$$ in an efficient way, then the problem is solved. However, it seems not so straightforward.
Section 2. The Trick for the Prefix Sum by Generating Functions
For convenience, we denote $$$$$$ S(a) = \frac {P(1)} {Q(1)} + \frac {P(2)} {Q(2)} + \dots + \frac {P(1)P(2) \dots P(a)} {Q(1)Q(2) \dots Q(a)}, $$$$$$ where $$$P(x)$$$ and $$$Q(x)$$$ are linear functions (polynomials of degree $$$1$$$). To compute $$$S(a)$$$, we use the trick mentioned in min_25's blog. However, the trick we adopt here can deal with more general cases.
Let $$$N > a$$$ be a large enough positive integer. Then we have $$$$$$ \begin{aligned} S(a) = \frac 1 {Q(1)Q(2) \dots Q(N)} \Big[ & P(1)Q(2)Q(3)\dots Q(a)Q(a+1) \dots Q(N) + \\ & P(1)P(2)Q(3)\dots Q(a)Q(a+1) \dots Q(N) + \\ & \dots + \\ & P(1)P(2)P(3)\dots P(a)Q(a+1) \dots Q(N) \Big]. \end{aligned} $$$$$$ We only have to compute the inner sum $$$s(a) = S(a)Q(1)Q(2) \dots Q(N)$$$. We now use the SQRT trick here. Let $$$B \approx \sqrt N$$$ be the block size. Let $$$$$$ \begin{aligned} R(z) = & P(z+1)Q(z+2)Q(z+3) \dots Q(z+B) + \\ & P(z+1)P(z+2)Q(z+3) \dots Q(z+B) + \\ & \dots + \\ & P(z+1)P(z+2)P(z+3) \dots P(z+B). \end{aligned} $$$$$$ Then we will see that $$$$$$ \begin{aligned} s(a) = & \sum_{k=0}^{\lfloor{a/B}\rfloor} P(1)\dots P(kB) R(kB) Q((k+1)B+1) \dots Q(N) \\ & + \sum_{i=(\lfloor{a/B}\rfloor+1)B+1}^a P(1)\dots P(i) Q(i+1) \dots Q(N). \end{aligned} $$$$$$ It is more convenient to let $$$\mathfrak{P}(k) = P(1)\dots P(kB)$$$ and $$$\mathfrak{Q}(k) = Q(kB+1) \dots Q(N)$$$, then the main summation in $$$s(a)$$$ is written as $$$$$$ \sum_{k=0}^{\lfloor{a/B}\rfloor} \mathfrak{P}(k) R(kB) \mathfrak{Q}(k+1). $$$$$$ Now, we have three values needed, which are $$$R(z), \mathfrak{P}(k), \mathfrak{Q}(k)$$$.
Compute $$$R(z)$$$
We deal with $$$R(z)$$$ first. We note that $$$R(z)$$$ is a polynomial of degree $$$B$$$ and we want to compute $$$R(kB)$$$ for every $$$0 \leq k \leq N/B$$$. If we can obtain the values of $$$R(0), R(1), \dots, R(B)$$$, then by Lagrange interpolation, we can recover the polynomial $$$R(z)$$$. And then, if we know the exact polynomial $$$R(z)$$$, then by multi-point evaluation trick, we can compute the values of $$$R(kB)$$$ for every $$$k \leq N/B$$$. The total time complexity would be $$$O(\sqrt{N} \log^2 N)$$$.
We can compute $$$R(0)$$$ in $$$O(B)$$$ time. And then, we note that $$$$$$ R(z) = \frac {Q(z+B)} {P(z)} R(z-1) - Q(z+1) \dots Q(z+B) + P(z+1) \dots P(z+B). $$$$$$ Therefore, $$$R(1), \dots, R(B)$$$ can be computed in $$$O(B)$$$ time if we simultaneously maintain the values of block product of $$$P(\cdot)$$$ and $$$Q(\cdot)$$$ in the above identity.
Bonus: A "log" in the complexity may be eliminated by some trick similar to min_25's blog?
Compute $$$\mathfrak{P}(k)$$$
This subproblem is very similar to compute the factorial-like product, which is originally mentioned in min_25's blog. To solve it, we define $$$$$$ f_P(z) = P(z+1)P(z+2) \dots P(z+B). $$$$$$ Then $$$\mathfrak{P}(k) = f_P(0) f_P(B) \dots f_P((k-1)B)$$$. What we need is to compute $$$f_P(kB)$$$ fast enough. We also note that $$$f_P(z)$$$ is a polynomial of degree $$$B$$$. We can easily find the values of $$$f_P(0), f_P(1), \dots, f_P(B)$$$ in $$$O(B)$$$ time. Therefore, the values of $$$f_P(kB)$$$ can be obtained by Lagrange interpolation and multi-point evaluation. In particular, we are able to compute $$$P(1)P(2)\dots P(N)$$$ in this way, which is considered to be a general factorial.
compute $$$\mathfrak{Q}(k)$$$
This subproblem is more tricky if we are able to compute $$$\mathfrak{P}(k)$$$ fast. Similarly, we define $$$f_Q(z)$$$ as what we did when computing $$$\mathfrak{P}(k)$$$. Let $$$\mathfrak{Q}'(k)$$$ be the prefix block products of $$$Q(x)$$$ as an analog for $$$\mathfrak{P}(k)$$$ and $$$P(x)$$$. It can be easily seen that $$$$$$ \mathfrak{Q}(k) = \frac {Q(1)Q(2)\dots Q(N)} {\mathfrak{Q}'(k)}. $$$$$$
Having computed all values we need in the previous several sections, we are able to compute the value of a single $$$f(a)$$$ in $$$O(\sqrt N)$$$ time. Therefore, the overall complexity is $$$O(\sqrt N \log^2 N + m \sqrt N)$$$.
Bonus: Saving $$$\log n$$$ in Complexity
Moreover, we can save $$$\log n$$$ in complexity with an improved interpolation method for shifting evaluation values. See the paper "A. Bostan, P. Gaudry, and E. Schost, Linear recurrences with polynomial coefficients and application to integer factorization and Cartier-Manin operator. SIAM Journal of Computing, 36(6): 1777–1806. 2007".
Shifting evaluation values
Suppose given are a polynomial $$$f(x)$$$ of degree $$$d$$$ by a point-value representation $$$f(0), f(1), \dots, f(d)$$$ and a shift $$$\delta$$$. It is asked to compute an alternative point-value representation $$$f(\delta), f(\delta+1), \dots, f(\delta+d)$$$. A straightforward way to solve this problem is to directly make use of Lagrange interpolation and multi-point evaluation, which yields an $$$O(d \log^2 d)$$$ arithmetic complexity algorithm. In fact, we can do it better in $$$O(d \log d)$$$ time.
Recall the formula of Lagrange interpolation that given $$$n$$$ points $$$(x_i, y_i)$$$ where $$$y_i = f(x_i)$$$, then the unique polynomial of degree $$$(n-1)$$$ is $$$$$$ f(x) = \sum_{i=1}^n y_{i} \prod_{j \neq i} \frac {x-x_j} {x_i-x_j}. $$$$$$ Rewrite this formula in our case that $$$f(0), f(1), \dots, f(d)$$$ are given, then $$$$$$ f(x) = \sum_{i=0}^d f(i) \prod_{j \neq i} \frac {x-j} {i-j}. $$$$$$ We assume that $$$\delta, \delta+1, \dots, \delta+d$$$ do not overlap with $$$0, 1, \dots, d$$$, then for $$$0 \leq k \leq d$$$, we have $$$$$$ \begin{aligned} f(\delta+k) & = \sum_{i=0}^d f(i) \prod_{j \neq i} \frac {\delta+k-j} {i-j} \\ & = \left( \prod_{j=0}^d (\delta+k-j) \right) \left( \sum_{i=0}^d \frac {(-1)^{d-i} f(i)} {i! (d-i)!} \cdot \frac {1} {\delta+k-i} \right). \end{aligned} $$$$$$ The summation can be computed by FFT. Consider two polynomials: $$$$$$ A(x) = \sum_{i=0}^d \frac {(-1)^{d-i} f(i)} {i! (d-i)!} x^i, $$$$$$ $$$$$$ B(x) = \sum_{i=0}^{2d} \frac {1} {\delta+i-d} x^i. $$$$$$ We have $$$$$$ f(\delta+k) = \left( \prod_{j=0}^d (\delta+k-j) \right) [x^{d+k}] A(x)B(x). $$$$$$
Therefore, $$$f(\delta), f(\delta+1), \dots, f(\delta+d)$$$ can be computed in $$$O(d \log d)$$$ time.
Arithmetic sequence shifting
Now we consider a modified version: given a polynomial of degree $$$d$$$ by a point-value representation $$$f(0), f(s), f(2s), \dots, f(ds)$$$ and a shift $$$\delta$$$, compute $$$f(\delta), f(\delta+s), f(\delta+2s), \dots, f(\delta+ds)$$$.
The trick is to let $$$g(x) = f(xs)$$$, which is also a polynomial of degree $$$d$$$. Then the problem becomes: given $$$g(0), g(1), \dots, g(s)$$$ and a shift $$$\delta' = \delta/s$$$, compute $$$g(\delta'), g(\delta'+1), \dots, g(\delta'+d)$$$.
Divide and conquer
Let's consider how to compute $$$R(z)$$$. For convenience, we denote that $$$$$$ \begin{aligned} R_B(z) = & P(z+1)Q(z+2)Q(z+3) \dots Q(z+B) + \\ & P(z+1)P(z+2)Q(z+3) \dots Q(z+B) + \\ & \dots + \\ & P(z+1)P(z+2)P(z+3) \dots P(z+B), \end{aligned} $$$$$$ $$$$$$ P_B(z) = P(z+1)P(z+2) \dots P(z+B), $$$$$$ $$$$$$ Q_B(z) = Q(z+1)Q(z+2) \dots Q(z+B), $$$$$$ where $$$B$$$ indicates the number of elements in the product. We need to compute the values of $$$$$$ \begin{matrix} R_B(0), & R_B(s), & \dots, & R_B(sB), \\ P_B(0), & P_B(s), & \dots, & P_B(sB), \\ Q_B(0), & Q_B(s), & \dots, & Q_B(sB). \end{matrix} $$$$$$
To compute these values recurrently, we should notice that $$$$$$ R_{2B}(z) = R_B(z) Q_B(z+B) + P_B(z) R_B(z+B), $$$$$$ $$$$$$ P_{2B}(z) = P_B(z) P_B(z+B), $$$$$$ $$$$$$ Q_{2B}(z) = Q_B(z) Q_B(z+B). $$$$$$ We choose $$$P_{2B}(z)$$$ as an example. Here we need to compute $$$$$$ \begin{aligned} P_{2B}(0) & = P_B(0) P_B(B), \\ P_{2B}(s) & = P_B(s) P_B(s+B), \\ \dots, \\ P_{2B}(sB) & = P_B(sB) P_B(sB+B), \\ P_{2B}(sB+1) & = P_B(sB+1) P_B(sB+B+1), \\ \dots, \\ P_{2B}(2sB) & = P_B(2sB) P_B(2sB+B). \end{aligned} $$$$$$ There are $$$4B$$$ values of $$$P_B(z)$$$ that are needed, which are $$$$$$ \begin{matrix} P_B(0), & P_B(s), & \dots, & P_B(sB), \\ P_B(B), & P_B(B+s), & \dots, & P_B(B+sB), \\ P_B(sB), & P_B(sB+s), & \dots, & P_B(sB+sB), \\ P_B(B+sB), & P_B(B+sB+s), & \dots, & P_B(B+sB+sB). \end{matrix} $$$$$$ The first line is our target and the rest three lines can be obtained by shifting the first line by $$$B$$$, $$$sB$$$ and $$$B+sB$$$, respectively. Let $$$T(d)$$$ be the complexity where $$$d$$$ is the degree of our polynomials, then $$$$$$ T(d) = T(d/2) + O(d \log d). $$$$$$ Hence, $$$T(d) = O(d \log d)$$$.
In our case, $$$d = O(\sqrt N)$$$, then the complexity is reduced to $$$O(\sqrt N \log N + m \sqrt N)$$$.
Conclusion
We use the trick, SQRT decomposition, to split summations into blocks, and then use Lagrange interpolation and multi-point evaluation to compute the value in each block. We obtain an algorithm with complexity $$$O(\sqrt N \log^2 N + m \sqrt N)$$$.
A more subtle trick to save a $$$\log N$$$ is obtained by the shifting technique, which yields a faster algorithm with complexity $$$O(\sqrt N \log N + m \sqrt N)$$$. The constant factor of Lagrange interpolation and multi-point evaluation is very large, so this approach to save a $$$\log N$$$ brings a significant improvement.
However, we allow suboptimal solutions that directly use Lagrange interpolation and multi-point evaluation to get AC.
Credits
Some related problems:
1. Stopping time: 1349D, 850F, 1025G.
2. Generating function: SPOJ FACTMODP, CodeChef QPOLYSUM, 1349F2.