A16878. Fair and Square
编程题
入门
知识点
题目描述
树是无向连通且无环的图。
给定一棵包含 $n$ 个顶点的树。每个顶点 $i$ 上写有一个整数 $a_i$。
对于任意两个顶点 $u$ 和 $v$（$u \ne v$），定义 $p(u, v)$ 为从 $u$ 到 $v$ 的唯一简单路径$^{\text{∗}}$ 上所有顶点所写数值的乘积。
一个由三个互不相同的顶点组成的无序三元组 $\{u, v, w\}$ 被称为“好”的，当且仅当：$p(u,v)\cdot p(v,w)\cdot p(w,u)$ 是一个完全平方数。
求给定树中“好”的无序三元组的个数。
$^{\text{∗}}$ 从顶点 $u$ 到顶点 $v$ 的简单路径是指一个顶点序列 $u = x_0, x_1, \ldots, x_k = v$，其中所有顶点互不相同，且对每个 $1 \le i \le k$，顶点 $x_{i-1}$ 与 $x_i$ 之间均存在一条边。
给定一棵包含 $n$ 个顶点的树。每个顶点 $i$ 上写有一个整数 $a_i$。
对于任意两个顶点 $u$ 和 $v$（$u \ne v$），定义 $p(u, v)$ 为从 $u$ 到 $v$ 的唯一简单路径$^{\text{∗}}$ 上所有顶点所写数值的乘积。
一个由三个互不相同的顶点组成的无序三元组 $\{u, v, w\}$ 被称为“好”的，当且仅当：$p(u,v)\cdot p(v,w)\cdot p(w,u)$ 是一个完全平方数。
求给定树中“好”的无序三元组的个数。
$^{\text{∗}}$ 从顶点 $u$ 到顶点 $v$ 的简单路径是指一个顶点序列 $u = x_0, x_1, \ldots, x_k = v$，其中所有顶点互不相同，且对每个 $1 \le i \le k$，顶点 $x_{i-1}$ 与 $x_i$ 之间均存在一条边。
输入格式
第一行包含一个整数 $t$（$1 \le t \le 10^4$）——测试用例的数量。随后是每个测试用例的描述。
每个测试用例以一个整数 $n$（$3 \le n \le 2\cdot 10^5$）开始——顶点的数量。
第二行包含 $n$ 个整数 $a_1,a_2,\dots,a_n$（$1 \le a_i \le 10^6$）——写在各顶点上的整数值。
接下来的 $n-1$ 行中，每行包含两个整数 $u,v$（$1 \le u,v \le n$），表示树的一条边。保证这些边构成一棵树。
保证所有测试用例的 $n$ 值之和不超过 $2\cdot 10^5$。
每个测试用例以一个整数 $n$（$3 \le n \le 2\cdot 10^5$）开始——顶点的数量。
第二行包含 $n$ 个整数 $a_1,a_2,\dots,a_n$（$1 \le a_i \le 10^6$）——写在各顶点上的整数值。
接下来的 $n-1$ 行中，每行包含两个整数 $u,v$（$1 \le u,v \le n$），表示树的一条边。保证这些边构成一棵树。
保证所有测试用例的 $n$ 值之和不超过 $2\cdot 10^5$。
输出格式
对于每个测试用例，输出树中好三元组的数量。
输入输出样例
输入 #1
4 5 1 1 1 1 1 1 2 2 3 2 4 4 5 10 1 2 3 4 5 6 7 8 9 10 1 3 2 6 6 7 5 4 8 3 3 4 4 6 9 1 10 2 6 12 6 3 18 9 2 3 4 4 5 2 6 6 1 4 2 8 3 16 9 1 8 16 4 9 2 1 3 1 4 3 3 5 6 3 4 7 8 1
输出 #1
10 48 0 40
说明/提示
对于第一个测试用例，所有由三个互异顶点构成的无序三元组都是“好”的：
1. $\{1, 2, 3\}$
2. $\{1, 2, 4\}$
3. $\{1, 2, 5\}$
4. $\{1, 3, 4\}$
5. $\{1, 3, 5\}$
6. $\{1, 4, 5\}$
7. $\{2, 3, 4\}$
8. $\{2, 3, 5\}$
9. $\{2, 4, 5\}$
10. $\{3, 4, 5\}$
对于第二个测试用例，$\{2, 5, 8\}$ 是一个“好”的三元组。
1. $\{1, 2, 3\}$
2. $\{1, 2, 4\}$
3. $\{1, 2, 5\}$
4. $\{1, 3, 4\}$
5. $\{1, 3, 5\}$
6. $\{1, 4, 5\}$
7. $\{2, 3, 4\}$
8. $\{2, 3, 5\}$
9. $\{2, 4, 5\}$
10. $\{3, 4, 5\}$
对于第二个测试用例，$\{2, 5, 8\}$ 是一个“好”的三元组。
