Shohag Loves Pebae
Contest: CodeTON Round 9 (Div. 1 + Div. 2, Rated, Prizes!)
Difficulty rating: 3500
Tags: math, number theory

Shohag has a tree with $$$n$$$ nodes.

Pebae has an integer $$$m$$$. She wants to assign each node a value — an integer from $$$1$$$ to $$$m$$$. So she asks Shohag to count the number, modulo $$$998\,244\,353$$$, of assignments such that following conditions are satisfied:

- For each pair $$$1 \le u \lt v \le n$$$, the least common multiple (LCM) of the values of the nodes in the unique simple path from $$$u$$$ to $$$v$$$ is not divisible by the number of nodes in the path.
- The greatest common divisor (GCD) of the values of all nodes from $$$1$$$ to $$$n$$$ is $$$1$$$.

But this problem is too hard for Shohag to solve. As Shohag loves Pebae, he has to solve the problem. Please save Shohag!

Input format:
The first line contains two space-separated integers $$$n$$$ and $$$m$$$ ($$$2 \le n \le 10^6$$$, $$$1 \le m \le 10^{9}$$$).

Each of the next $$$n - 1$$$ lines contains two integers $$$u$$$ and $$$v$$$ ($$$1 \le u, v \le n$$$) indicating there is an edge between vertices $$$u$$$ and $$$v$$$. It is guaranteed that the given edges form a tree.

Output format:
Print a single integer — the number of valid ways to assign each vertex a value, modulo $$$998\,244\,353$$$.

Constraints:
Time limit: 5.0 s
Memory limit: 768.0 MB

Note:
In the first test case, the valid assignments are $$$[1, 1, 1, 1, 1, 1]$$$ and $$$[1, 1, 1, 1, 1, 5]$$$.

In the second test case, the valid assignments are $$$[1, 1]$$$, $$$[1, 3]$$$, $$$[1, 5]$$$, $$$[3, 1]$$$, $$$[3, 5]$$$, $$$[5, 1]$$$ and $$$[5, 3]$$$.

Example input:
6 6
1 2
2 3
3 4
4 5
3 6

Example output:
2

Example input:
2 5
1 2

Example output:
7

Editorial:
Let's say we assign $$$a_u$$$ to the node $$$u$$$. Let $$$h_u$$$ be the maximum length of a simple path that passes through $$$u$$$. Then a necessary condition is that $$$a_u$$$ can not be a multiple of any number $$$\le h_u$$$. Because if $$$a_u$$$ is a multiple of $$$k \le h_u$$$ and $$$v$$$ is a node such that the unique simple path from $$$u$$$ to $$$v$$$ has length $$$k$$$, then the LCM of the values of the nodes from $$$u$$$ to $$$v$$$ is a multiple of $$$k$$$, which is a contradiction.
The condition also means that $$$a_u$$$ can not be a multiple of any prime number $$$p \le h_u$$$.
Is this a sufficient condition? Yes, and the proof is also simple.
So now the problem is to count the number of assignments such that for each node $$$u$$$, $$$a_u$$$ is not a multiple of any prime number $$$p \le h_u$$$ and $$$\text{gcd}(a_1, a_2, \ldots, a_n) = 1$$$.
Let
• $$$f_{w, p}$$$ be the count of numbers from $$$1$$$ to $$$w$$$ that are not divisible by any prime $$$\le p$$$,
• $$$D$$$ be the diameter of the tree,
• A number $$$x$$$ is good if $$$x$$$ is not divisible by any prime $$$\le D$$$,
• $$$\mu(g)$$$ be the Mobius function,
• $$$\pi(x)$$$ be the number of primes $$$\le x$$$.
Then the answer to our problem is $$$\sum_{g = 1}^m \mu(g) \cdot [g \text{ is good}] \cdot \prod_{i = 1}^n f_{\lfloor \frac{m}{g} \rfloor, h_i}$$$.
As $$$\lfloor \frac{m}{g} \rfloor$$$ is a non-decreasing function and has at most $$$2 \sqrt{m}$$$ distinct values, we can iterate over $$$\lfloor \frac{m}{g} \rfloor$$$ and calculate range sums of $$$\mu(g) \cdot [g \text{ is good}]$$$.
For calculating prefix sums of a multiplicative function (like $$$\mu(g)$$$), it's a standard task and can be solved using Dirichlet convolution, Min25 sieve or multiple other methods.
Here, we need a slight variant of the method as we need the prefix sums of $$$\mu(g) \cdot [g \text{ is good}]$$$. This can be achieved using Dirichlet convolution in $$$\mathcal{O}(m^{2 / 3})$$$ if we just imagine the prime numbers $$$\le D$$$ do not exist in the number system. Refer to my code for more details.
But for each fixed $$$\lfloor \frac{m}{g} \rfloor$$$, how do we calculate $$$\prod_{i = 1}^n f_{\lfloor \frac{m}{g} \rfloor, h_i}$$$ fast enough? Trivially doing it will make the total complexity around $$$\mathcal{O}(n \sqrt{m})$$$ which is too slow.
The key observation is to not forget that the values of $$$h_i$$$ are not random, they are the maximum length of a simple path that passes through the node $$$i$$$. So $$$h_i \ge \lceil \frac{D}{2} \rceil$$$ for all $$$i$$$ because from each node, the endpoints of the diameter are at least $$$\lceil \frac{D}{2} \rceil$$$ away.
So now consider two cases:
Case 1: $$$D > 2 \sqrt{m}$$$
In this case, all $$$h_i \ge \lceil \frac{D}{2} \rceil \ge \sqrt{m}$$$ for all $$$i$$$. So only primes or $$$1$$$ are the good numbers. So instead of going with the mobius route, we can just directly solve it by calculating the total number of ways and subtracting the number of ways where the gcd is a prime.
We can calculate the total number of ways by first calculating the number of primes $$$\le m$$$ and then $$$f_{m, h_i}$$$ is just $$$\pi(m) - \pi(h_i) + 1$$$.
And the number of ways where the gcd is a prime is just $$$1$$$ for all primes $$$> D$$$ and $$$0$$$ otherwise.
Counting primes under $$$m$$$ is also a standard task and can be done in $$$\mathcal{O}(m^{2 / 3} \log m)$$$ or faster.
Case 2: $$$D \le 2 \sqrt{m}$$$
We can convert each $$$h_i$$$ to the maximum prime $$$\le h_i$$$ and then group $$$h_i$$$ by their values. Then the maximum number of groups will be $$$\mathcal{O}(\pi(\sqrt{m}))$$$. So for each fixed $$$k = \lfloor \frac{m}{g} \rfloor$$$, if the sum of the mobius function in the range $$$(\lfloor \frac{m}{k + 1} \rfloor, \lfloor \frac{m}{k} \rfloor]$$$ is non-zero (keep in mind that when all numbers in the range are bad numbers, then the sum will definitely be $$$0$$$), then we can calculate the product of $$$f_{k, h_i}$$$ directly. Then the upper bound of the complexity will be around $$$\mathcal{O}(\frac{m}{\log^2 m} \cdot \log \left( \frac{n}{\pi(\sqrt[3]{m})} \right))$$$. The proof will be added later. This works fast enough.