Fair and Square

The condition looks number-theory-heavy, but the tree structure nukes almost all of it.

The Square Part

For a product to be a perfect square, every prime must have an even exponent. So we only care about exponents modulo $2$.

Now take any three distinct vertices $u,v,w$. In a tree, the three paths between them meet at exactly one vertex. Call it $m$. This is the usual tree median: the unique common vertex of the paths $u\leftrightarrow v$, $v\leftrightarrow w$, and $w\leftrightarrow u$.

Consider the product

$$p(u,v)\cdot p(v,w)\cdot p(w,u).$$

Look at how many times each vertex value is multiplied in.

Every vertex on one of the three arms from $m$ to $u,v,w$ appears on exactly two of the three pair paths. That includes the endpoints $u,v,w$ unless one of them is $m$.

The median vertex $m$ appears on all three paths.

So modulo parity:

• every non-median vertex contributes $0$ times modulo $2$,
• the median contributes $1$ time modulo $2$.

Therefore the whole product has the same square/parity behavior as just $a_m$.

So:

$$\{u,v,w\}\text{ is good} \iff a_m\text{ is a perfect square}.$$

That’s the key. The products were bait. Pretty rude bait, honestly.

So now the problem is:

Count unordered triples of vertices whose tree median has a perfect-square value.

Counting Triples With Fixed Median

Fix a vertex $x$. Suppose we remove $x$ from the tree. The remaining graph splits into connected components with sizes

$$s_1,s_2,\dots,s_k.$$

These components are exactly the neighbor-side subtrees around $x$.

A triple has median $x$ in two possible ways.

Case 1: The triple contains $x$

Then the other two vertices must lie in two different components around $x$.

If both were in the same component, the path between them would not need to pass through $x$, so $x$ would not be the median.

The count is

$$\sum_{i<j} s_i s_j.$$

Case 2: The triple does not contain $x$

Then the three vertices must lie in three different components around $x$.

The count is

$$\sum_{i<j<l} s_i s_j s_l.$$

So if $a_x$ is a perfect square, vertex $x$ contributes

$$\sum_{i<j} s_i s_j + \sum_{i<j<l} s_i s_j s_l.$$

If $a_x$ is not a perfect square, it contributes $0$.

Getting Component Sizes Fast

Root the tree anywhere, say at $1$.

Compute subtree sizes $sub[x]$.

For a vertex $x$, each child $c$ contributes one component of size $sub[c]$ after removing $x$.

If $x$ has a parent, the parent-side component has size

$$n - sub[x].$$

So we can build the list of component sizes around every vertex in total $O(n)$.

Computing the Sums

For sizes $s_1,\dots,s_k$:

Pair sum:

$$\sum_{i<j}s_is_j$$

Triple sum:

$$\sum_{i<j<l}s_is_js_l$$

We can compute both in one pass:

Maintain:

• $sum_1 =$ sum of previous sizes,
• $sum_2 =$ sum of products of pairs among previous sizes,
• $ways =$ total contribution.

When adding a new size $s$:

• new pairs formed: $sum_1\cdot s$,
• new triples formed: $sum_2\cdot s$.

So add both to the answer, then update:

$$sum_2 \mathrel{+}= sum_1s,$$ $$sum_1 \mathrel{+}= s.$$

Only do this for vertices whose $a_x$ is a perfect square.

Perfect Square Check

Since $a_i\le 10^6$, just compute

$$r=\lfloor \sqrt{a_i}\rfloor$$

and check $r^2=a_i$ or $(r+1)^2=a_i$ to dodge floating-point nonsense.

Complexity

Each test case is linear:

$$O(n)$$

The total $n$ over all tests is at most $2\cdot 10^5$, so this easily fits. No prime factorization needed. The tags tried to sell you number theory; the tree median did the actual work.