Strange Function
Contest: VK Cup 2019-2020 - Elimination Round (Engine)
Difficulty rating: 2900
Tags: dp

Let's define the function $$$f$$$ of multiset $$$a$$$ as the multiset of number of occurences of every number, that is present in $$$a$$$.

E.g., $$$f(\{5, 5, 1, 2, 5, 2, 3, 3, 9, 5\}) = \{1, 1, 2, 2, 4\}$$$.

Let's define $$$f^k(a)$$$, as applying $$$f$$$ to array $$$a$$$ $$$k$$$ times: $$$f^k(a) = f(f^{k-1}(a)), f^0(a) = a$$$.

E.g., $$$f^2(\{5, 5, 1, 2, 5, 2, 3, 3, 9, 5\}) = \{1, 2, 2\}$$$.

You are given integers $$$n, k$$$ and you are asked how many different values the function $$$f^k(a)$$$ can have, where $$$a$$$ is arbitrary non-empty array with numbers of size no more than $$$n$$$. Print the answer modulo $$$998\,244\,353$$$.

Input format:
The first and only line of input consists of two integers $$$n, k$$$ ($$$1 \le n, k \le 2020$$$).

Output format:
Print one number — the number of different values of function $$$f^k(a)$$$ on all possible non-empty arrays with no more than $$$n$$$ elements modulo $$$998\,244\,353$$$.

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

Example input:
3 1

Example output:
6

Example input:
5 6

Example output:
1

Editorial:
The solution of the task consists of three cases:
• $$$k = 1$$$. For fixed $$$n$$$ $$$f(a)$$$ can be equal to any partition of $$$n$$$. We need to count the number of arrays $$$b_1, b_2, \ldots, b_m$$$, such that $$$b_1 \ge b_2 \ge \ldots \ge b_m$$$ and $$$\sum \limits_{i=1}^m b_i \le n$$$. This can be done by simple dp in $$$\mathcal{O}(n^2)$$$ (or even faster, much faster).
• $$$k = 2$$$. When the array $$$b_1 \ge b_2 \ge \ldots \ge b_m$$$ can be equal to value of $$$f^2(a)$$$ for some $$$|a| \le n$$$? When there exists some array $$$c_1, \ldots, c_l$$$ such that $$$\sum \limits_{i=1}^l c_i \le n$$$, and $$$f(c) = b$$$. The values of $$$b$$$ are the numbers of occurences of numbers in $$$c$$$, so we need to minimize $$$\sum \limits_{i=1}^m b_i v_i$$$, where $$$v_i$$$ — the unique numbers in $$$c$$$. To minimize this sum we should take $$$v_1=1, \ldots, v_m=m$$$, so we need $$$\sum \limits_{i=1}^m b_i i \le n$$$. This can be done by simple dp: $$$dp[val][j][sum]$$$ — the number of prefixes of $$$b$$$ such that we already took $$$j$$$ elements to $$$b$$$, all elements on prefix are greater than or equal to $$$val$$$, and the $$$\sum \limits_{i=1}^j b_i i = sum$$$.This dp can look like it is $$$\mathcal{O}(n^3)$$$, but it is actually $$$\mathcal{O}(n^2 \log n)$$$, because there is a limitation $$$val \cdot j \le n$$$, and there are $$$\mathcal{O}(n \log n)$$$ such pairs. There is also a subquadratic solution.
• $$$k \ge 3$$$. We can notice that in the array $$$f^2(a)$$$ there are at most $$$\mathcal{O}(\sqrt{2n})$$$ elements. We can use this fact to bruteforce all possible answers — candidates for the answer are the partitions of numbers not exceeding $$$\sqrt{2n}=64$$$, there are few millions of them. How to check if the array $$$b_1 \ge b_2 \ge \ldots \ge b_m$$$ can be the falue of $$$f^k(a)$$$? It happens that we can make $$$k-2$$$ iterations of the unfolding algorithm from case $$$k=2$$$ and get the «minimal» possible array $$$a$$$, and check, if it contains no more than $$$n$$$ elements. This part works in $$$\mathcal{O}(\mathcal{P}(\sqrt{2n}))$$$.