题目描述

EJOI 期间，Anton 在 Kaunas 散步时发现了 $N$ 个参赛者聚集点，编号为 $0$ 到 $N-1$。每个地点恰好由一支队伍的成员占据，队伍编号为 $0$ 到 $K-1$。同一支队伍的成员可以占据任意多个地点，也可能有队伍不占据任何地点。

这些地点由 $N-1$ 条双向道路连接，并且任意两个地点之间都恰好存在一条简单路径，因此它们构成一棵树。简单路径是一个由互不相同的地点组成的序列，其中每两个相邻地点之间都有道路。路径长度是它使用的道路数，也就是经过的地点数减 $1$。

Anton 想沿一条简单路径散步，并经过尽可能多的地点。如果一条简单路径的长度在树中所有简单路径里最大，则称它是有趣路径。一条路径的团队丰富度是 Anton 沿途遇到的不同队伍数量。

请计算所有不同有趣路径的团队丰富度之和。当且仅当两条有趣路径经过的地点集合完全相同时，它们才被视为同一条路径。特别地，反向经过一条路径不会产生新的路径。

实现细节

你需要实现以下函数：

long long teamfulness(int N, int K, std::vector<int> a,
                      std::vector<int> u, std::vector<int> v)

• $N$：地点数；
• $K$：队伍数；
• $a$：长度为 $N$ 的数组，其中 $a_i$ 表示占据地点 $i$ 的队伍；
• $u,v$：长度为 $N-1$ 的数组，其中 $u_i$ 和 $v_i$ 是第 $i$ 条道路连接的两个地点。

每个测试中，该函数恰好调用一次，并且必须返回所有有趣路径的团队丰富度之和。

输入格式

输入格式：

• 第 $1$ 行：两个整数 $N$ 和 $K$；
• 第 $2$ 行：$N$ 个整数 $a_0,a_1,\ldots,a_{N-1}$；
• 第 $3+i$ 行：两个整数 $u_i$ 和 $v_i$，即第 $i$ 条道路的两个端点。

输出格式

输出格式：

• 第 $1$ 行：函数的返回值。

6 3
1 0 0 1 2 1
0 1
0 2
0 3
0 4
0 5

21

7 1
0 0 0 0 0 0 0
0 1
0 2
1 3
1 4
2 5
2 6

4

6 3
0 1 2 0 1 2
0 1
1 2
2 3
1 4
2 5

11

提示

样例 1 解释

简单路径的最大长度为 $2$，因此有趣路径包含 $2$ 条道路和 $3$ 个地点。团队丰富度为 $3$ 的有趣路径有 $2$ 条，团队丰富度为 $2$ 的有 $7$ 条，团队丰富度为 $1$ 的有 $1$ 条，总和为 $21$。

样例 2 解释

图中队伍 $0$ 用黄色表示：

:::align{center} :::

由于所有地点都属于唯一的一支队伍，每条路径的团队丰富度均为 $1$。长度为 $4$ 的有趣路径共有 $4$ 条，因此总和为 $4$。

样例 3 解释

图中队伍 $0$ 为黄色，队伍 $1$ 为绿色，队伍 $2$ 为红色：

:::align{center} :::

有趣路径的长度为 $3$。共有 $4$ 条有趣路径，其中 $3$ 条的团队丰富度为 $3$，另 $1$ 条为 $2$，总和为 $11$。

限制

• $3\le N\le 10^6$
• $1\le K\le N$
• 对每个 $0\le i<N$，均有 $0\le a_i<K$
• 对每个 $0\le i<N-1$，均有 $0\le u_i,v_i<N$

子任务