Electrical Efficiency
Contest: COMPFEST 14 - Preliminary Online Mirror (Unrated, ICPC Rules, Teams Preferred)
Difficulty rating: 2500
Tags: combinatorics, data structures, dp, math, number theory, trees

In the country of Dengkleknesia, there are $$$N$$$ factories numbered from $$$1$$$ to $$$N$$$. Factory $$$i$$$ has an electrical coefficient of $$$A_i$$$. There are also $$$N-1$$$ power lines with the $$$j$$$-th power line connecting factory $$$U_j$$$ and factory $$$V_j$$$. It can be guaranteed that each factory in Dengkleknesia is connected to all other factories in Dengkleknesia through one or more power lines. In other words, the collection of factories forms a tree. Each pair of different factories in Dengkleknesia can use one or more existing power lines to transfer electricity to each other. However, each power line needs to be turned on first so that electricity can pass through it.

Define $$$f(x, y, z)$$$ as the minimum number of power lines that need to be turned on so that factory $$$x$$$ can make electrical transfers to factory $$$y$$$ and factory $$$z$$$. Also define $$$g(x, y, z)$$$ as the number of distinct prime factors of $$$\text{GCD}(A_x, A_y, A_z)$$$.

To measure the electrical efficiency, you must find the sum of $$$f(x, y, z) \times g(x, y, z)$$$ for all combinations of $$$(x, y, z)$$$ such that $$$1 \leq x < y < z \leq N$$$. Because the answer can be very large, you just need to output the answer modulo $$$998\,244\,353$$$.

Note: $$$\text{GCD}(k_1, k_2, k_3)$$$ is the greatest common divisor of $$$k_1$$$, $$$k_2$$$, and $$$k_3$$$, which is the biggest integer that simultaneously divides $$$k_1$$$, $$$k_2$$$, and $$$k_3$$$.

Input format:
The first line contains a single integer $$$N$$$ ($$$1 \le N \le 2 \cdot 10^5$$$) — the number of factories in Dengkleknesia.

The second line contains $$$N$$$ integers $$$A_1, A_2, \dots, A_N$$$ ($$$1 \leq A_i \leq 2 \cdot 10^5$$$) — the electrical coefficients of the factories.

The $$$j$$$-th of the next $$$N-1$$$ lines contains two integers $$$U_j$$$ and $$$V_j$$$ ($$$1 \le U_j, V_j \le N$$$) — a power line that connects cities $$$U_j$$$ and $$$V_j$$$. The collection of factories forms a tree.

Output format:
An integer representing the sum of $$$f(x, y, z) \times g(x, y, z)$$$ for all combinations of $$$(x, y, z)$$$ such that $$$1 \leq x < y < z \leq N$$$, modulo $$$998\,244\,353$$$

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

Note:
In the first example, the only $$$(x, y, z)$$$ possible is $$$(1, 2, 3)$$$. Because $$$\text{GCD}(A_1, A_2, A_3) = \text{GCD}(1, 2, 3) = 1$$$ has $$$0$$$ distinct prime factors, therefore $$$f(x, y, z) \times g(x, y, z) = 2 \times 0 = 0$$$.

In the second example, all triples $$$(x, y, z)$$$ that satisfy the condition are as follows:

- $$$(1, 2, 3) \rightarrow f(1, 2, 3) \times g(1, 2, 3) = 2 \times 1 = 2$$$
- $$$(1, 2, 4) \rightarrow f(1, 2, 4) \times g(1, 2, 4) = 2 \times 1 = 2$$$
- $$$(1, 3, 4) \rightarrow f(1, 3, 4) \times g(1, 3, 4) = 3 \times 2 = 6$$$
- $$$(2, 3, 4) \rightarrow f(2, 3, 4) \times g(2, 3, 4) = 2 \times 1 = 2$$$

So the electrical efficiency is $$$2 + 2 + 6 + 2 = 12$$$.

Example input:
3
1 2 3
1 2
2 3

Example output:
0

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

Example output:
12

Editorial:
Note that $$$\sum_x \sum_y \sum_z f(x,y,z) \times g(x,y,z)$$$ can be alternatively expressed as the following:
• Iterate through every edge $$$e$$$ and every prime $$$p$$$.
• In each iteration, consider the set of vertices $$$S = \{x \mid A_x \text{ is divisible by } p\}$$$. Count how many triplets $$$a,b,c \in S$$$ such that each of the two components separated by $$$e$$$ contains at least one of $$$a,b,c$$$.
• Calculate the sum of those values for all possible $$$e$$$ and $$$p$$$.
Note that if we root the tree and consider the edge $$$(x,y)$$$ where $$$x$$$ is the parent of $$$y$$$ in the rooted tree, then the number of triplets we count in one iteration is equivalent to $$$\binom{|S|}{3} - \binom{s}{3} - \binom{|S|-s}{3}$$$, where $$$s$$$ is the number of vertices in $$$S$$$ in the subtree of $$$y$$$.
To calculate the answer, we can iterate through each prime $$$p$$$. For each prime, we only consider the vertices in $$$S$$$. We can build a sparse representation of a tree containing the vertices in $$$S$$$ (i.e. a virtual/auxiliary tree). More precisely, we only consider a set of vertices $$$S'$$$, where $$$x \in S'$$$ if and only if $$$x \in S$$$ or there exists $$$y,z \in S$$$ such that $$$\text{LCA}(y,z) = x$$$ (here, $$$\text{LCA}$$$ denotes the lowest common ancestor). It can be shown that $$$|S'| \leq 2|S|$$$. Each edge in the sparse tree connecting two vertices in $$$S'$$$ is a simple ancestor-descendant path consisting of one or more edges of the original tree.
Notice that the number of triplets we count for each edge in the original tree that is not used in the sparse tree is always $$$0$$$. Also, for the edges in the original tree that make up a single edge of the sparse tree, the number of triplets we count for each of them is the same. Therefore, using the sparse tree we built, we can run a simple dynamic programming to count the answer for a single prime $$$p$$$ in $$$O(|S|)$$$ time.
Notice that there are at most $$$\log A_x$$$ distinct prime factors of $$$A_x$$$, so the sum of $$$|S|$$$ for all primes is $$$O(N \log \max(A))$$$.
Time complexity: $$$O(N \log N \log \max(A))$$$ or $$$O(N (\log N + \log \max(A)))$$$