题目描述

While walking around Kaunas during EJOI, Anton discovered $N$ spots where participants hang out. The spots are numbered from $0$ to $N-1$. Each spot is occupied by members of exactly one team, and teams are numbered from $0$ to $K-1$. Members of the same team may occupy any number of spots, and some teams may occupy no spots.

The spots are connected by $N-1$ two-way roads such that there is exactly one simple path between any two spots; therefore, they form a tree. A simple path is a sequence of distinct spots in which every two consecutive spots are connected by a road. Its length is the number of roads it uses, one less than the number of spots it visits.

Anton wants to walk along a simple path and visit as many spots as possible. A simple path is interesting if its length is maximum among all simple paths in the tree. The teamfulness of a path is the number of distinct teams Anton encounters along it.

Find the sum of teamfulness over all different interesting paths. Two interesting paths are considered the same if and only if they visit exactly the same set of spots. In particular, traversing a path in the opposite direction does not create a different path.

Implementation details

Implement the following function:

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

• $N$: the number of spots;
• $K$: the number of teams;
• $a$: an array of $N$ integers, where $a_i$ is the team occupying spot $i$;
• $u,v$: arrays of $N-1$ integers, where $u_i$ and $v_i$ are the spots connected by the $i$-th road.

The function is called exactly once per test and must return the sum of teamfulness over all interesting paths.

输入格式

Input format:

• line $1$: $N$ and $K$;
• line $2$: $N$ integers $a_0,a_1,\ldots,a_{N-1}$;
• line $3+i$: two integers $u_i$ and $v_i$, the endpoints of the $i$-th road.

输出格式

Output format:

• line $1$: the value returned by the function.

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

提示

Explanation of example 1

The maximum length of a simple path is $2$, so interesting paths have $2$ roads and $3$ spots. There are two interesting paths with teamfulness $3$, seven with teamfulness $2$, and one with teamfulness $1$, giving a total of $21$.

Explanation of example 2

Team $0$ is shown in yellow:

:::align{center} :::

Every path has teamfulness $1$ because there is only one team. There are four interesting paths of length $4$, so the sum is $4$.

Explanation of example 3

Team $0$ is yellow, team $1$ is green, and team $2$ is red:

:::align{center} :::

Interesting paths have length $3$. There are four interesting paths: three have teamfulness $3$, and one has teamfulness $2$. Their total teamfulness is $11$.

Constraints

• $3\le N\le 10^6$
• $1\le K\le N$
• $0\le a_i<K$ for every $0\le i<N$
• $0\le u_i,v_i<N$ for every $0\le i<N-1$

Subtasks

Subtask	Points	$N$	$K$	Additional constraints
0		-		The examples.
1	4	$\le 10^6$	$\le N$	Every spot is directly connected to at most two other spots.
2	7			Some spot is directly connected to every other spot.
3	9	$\le 200$		-
4	10	$\le 2\cdot 10^3$
5		$\le 10^6$	$=1$
6	9		$\le 2$
7	11	$\le 2\cdot 10^5$	$\le 50$
8	12		$\le N$
9	13	$\le 10^6$		The length of an interesting path is odd.
10	15			-