(In Chinese: https://www.cnblogs.com/YunQianQwQ/articles/19080422)

We can do a simple DP to solve problem 1 in \(O(N)\).

For problem 2, we need to find

\[ \sum_{A x + B y + C z = N} \frac{N!}{(A x)!(B y)!(C z)!} \]

If \(A\times B \ge \sqrt{N}\), we can enumerate \(x,y\), then compute \(z\) and calculate the contribution. The time complexity is \(O(\frac{N^2}{AB})\).

If \(A\times B \le \sqrt{N}\), we can first perform a DP: \(\text{dp}[i][x][y]\) denotes the number of ways after considering the first \(i\) balls, where balls can only be placed in boxes \(A\) and \(B\), and the number of balls in the two boxes modulo \(A\) and modulo \(B\) are \(x\) and \(y\), respectively. The final answer is \(\sum_{z}\binom{N}{Cz}\times \text{dp}[N-Cz][0][0]\). The time complexity is \(O(N\times A\times B)\).

Let’s use the first algorithm when \(AB\ge \sqrt{N}\), and use the second algorithm when \(AB\le \sqrt{N}\), and the total time complexity is \(O(N^{1.5})\).

https://atcoder.jp/contests/abc422/submissions/69165517