For the given set of \(N\) candies, consider when Takahashi randomly choose any of the \(2^N\) subsets of them (including empty set) with equal probability. Let \(p(k,c)\) be the probability that the subset Takahashi chose has exactly \(k\) candies with exactly \(c\) different colors. Let \(C_k\) be the random variable for the number of different colors given that exactly \(k\) candies are chosen. Let \(f_k(x)=\sum_{c=0}^{k}p(k,c)x^c\) be the (unnormalized) probability-generating function (PGF) of \(C_k\).

The task is to find \(\mathrm{E}[C_k]\) for all \(k=1,2,\dots,N\). We don’t need to (and cannot) find every \(p(k,c)\), because the expectation can be expressed by \(f_k\) as follows (\(f'_k\) is the derivative of \(f_k\) with respect to \(x\)):

\[ \mathrm{E}[C_k]=\frac{\sum_{c=0}^{k}cp(k,c)}{\sum_{c=0}^{k}p(k,c)}=\frac{f'_k(1)}{f_k(1)} \]

We can find \(f_k(1)\) and \(f'_k(1)\) for all \(k=1,2,\dots,N\) efficiently with the following divide-and-conquer approach.

For any two sets of candies with no color appearing in both sets, every way of choosing \(k_1\) candies with \(c_1\) different colors from the first set, and \(k_2\) candies with \(c_2\) different colors from the second set, corresponds to a way of choosing \(k_1+k_2\) candies with \(c_1+c_2\) different colors from the union of the two sets. Thus, formally, let the two sets be \(S_1\) and \(S_2\) and let \(f_k^{(S_1)}\) and \(f_k^{(S_2)}\) be their PGFs respectively, the PGFs of their union set \(f_k^{(S_1\cup S_2)}(x)\) can be expressed in the form of convolution.

\[ f_k^{(S_1\cup S_2)}(x)=\sum_{k_1+k_2=k}f_{k_1}^{(S_1)}(x)f_{k_2}^{(S_2)}(x) \]

Specifically,

\[ f_k^{(S_1\cup S_2)}(1)=\sum_{k_1+k_2=k}f_{k_1}^{(S_1)}(1)f_{k_2}^{(S_2)}(1) \]

and, by the product rule of derivatives,

\[ f'^{(S_1\cup S_2)}_k(1)=\sum_{k_1+k_2=k}\left( f'^{(S_1)}_{k_1}(1)f_{k_2}^{(S_2)}(1)+f_{k_1}^{(S_1)}(1)f'^{(S_2)}_{k_2}(1) \right) \]

Thus, the \(f_k(1)\) and \(f'_k(1)\) values of two sets can be merged using FFT in \(O(N\log N)\). To efficiently merge many sets into one, we can always merge two smallest sets in each step. The overall complexity is \(O(N\log^2N)\). This is faster than the \(O(N\sqrt{N})\) official solution.

Lastly, the base cases are sets of candies with only one color. It can be easily derived that \(f_k(1)={N\choose k}/2^N\), \(f'_0(1)=0\), \(f'_k(1)=f_k(1) (k\geq 1)\) in such cases.

Implementation example: https://atcoder.jp/contests/abc215/submissions/25239792