PROBLEM LINK:
Author: Satyaki Upadhyay
Tester: Misha Chorniy
Editorialist: Animesh Fatehpuria
PROBLEM
Count number of semi-BSTs having exactly n nodes. A semi-BST is a tree with n nodes which satisfies the following conditions:
- Each node contains a distinct value from \{1, 2, 3, \dots, n\}.
- The root has at most two children.
- The right subtree of the root (if it exists) is a valid semi-BST.
- The left subtree of the root (if it exists) is any unordered rooted tree. That is, the children of a vertex are not ordered, and hence no concept of left or right child.
- The BST property for the root is maintained, i.e. all the values in the left subtree of the root are lesser than that of the root and all the values in the right subtree of the root are greater than that of the root.
Constraints: n \le 10^5, \text{test cases} \le 10^5.
Print the answer modulo 663224321 i.e. 2^{19} \times 5 \times 11 \times 23 + 1.
PREREQUISITES
Knowledge of “online” FFT trick.
EXPLANATION
We first solve the problem in O(n^2), and then optimize it using a trick.
An O(n^2) Solution
We use a dynamic program similar to “count number of BSTs having n nodes”. Define T(n) to be number of semi-BSTs of size n. To find a recurrence for T(n), we will iterate over the number of nodes in the left subtree. Note that fixing the number of nodes in the left subtree determines the value of the root. Suppose we say that the left subtree has i
nodes. Then, the number of labeled trees with i nodes is i^{i - 2} by Cayley’s Formula. Thus, the number of ways for us to choose a left subtree is i \times i^{i - 2} = i^{i - 1} since we have i choices for the root of the left subtree as well. By definition, the right subtree can be chosen in T(n - 1 - i) ways. Note that these two are independent, making our recurrence:
T(n) = \sum_{i = 0}^{n - 1} (i^{i - 1} \times T(n - 1 - i)) with base cases T(0) = T(1) = 1.
Clearly, there are O(n) states, and it takes O(n) to compute each state. Therefore, this solution is O(n^2).
Speedup using Online FFT
It turns out that we can speed up the above dynamic program using a trick known as online FFT. The special modulo would have given this hint as well! For reference, please solve this Codeforces problem. The editorial for this problem explains this trick!
Additionally, you can check out these brilliant slides written by Tanuj Khattar! Interestingly, his team TooWeakTooSlow were the only team to solve this problem in the onsite regional.
The final complexity would be O(n \log^2{n}).
AUTHOR’S AND TESTER’S SOLUTIONS:
Author’s solution can be found here.
Tester’s solution can be found here.