• 比赛：EJOI 2026 Day 2
• 时间限制：4 秒
• 内存限制：1024 MiB
• 题目类型：函数提交题

题目描述

Anton 在 Kaunas 参加 EJOI 时发现了 $N$ 个选手聚集地点，编号为 $0$ 到 $N-1$。

每个地点恰好由一个队伍的成员占据。队伍编号为 $0$ 到 $K-1$。同一个队伍可以占据任意多个地点，也可能有某个队伍没有占据任何地点。

这些地点之间有 $N-1$ 条双向道路，并且任意两个地点之间恰好存在一条简单路径，因此这些地点和道路构成一棵树。

一条简单路径的长度定义为它经过的道路数量，即经过地点数减一。

Anton 想沿着一条简单路径散步，并尽可能经过更多地点。

若一条简单路径的长度在整棵树的所有简单路径中达到最大值，则称其为一条有趣路径。也就是说，有趣路径就是树的直径路径。

一条路径的 teamfulness（队伍丰富度） 定义为：这条路径上出现了多少个不同的队伍。

你的任务是求：

所有不同有趣路径的 teamfulness 之和。

两条有趣路径当且仅当经过的地点集合完全相同时才被视作同一条路径。特别地，将同一条路径反向行走不会产生一条新的路径。

实现要求

实现：

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

其中：

• N：地点数量；
• K：队伍数量；
• a[i]：地点 $i$ 所属的队伍；
• 对每个 $0\le i<N-1$，u[i] 和 v[i] 表示第 $i$ 条道路连接的两个地点。

函数每个测试恰好调用一次，并返回所有有趣路径的 teamfulness 总和。

数据范围

• $3\le N\le10^6$；
• $1\le K\le N$；
• $0\le a_i<K$；
• $0\le u_i,v_i<N$。

样例 1

Sample grader 输入

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

输出

21

说明

最长简单路径长度为 $2$。

共有：

• 2 条有趣路径的 teamfulness 为 3；
• 7 条有趣路径的 teamfulness 为 2；
• 1 条有趣路径的 teamfulness 为 1。

因此总和为

$2\cdot3+7\cdot2+1=21$。

样例 2

Sample grader 输入

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

输出

4

只有一个队伍，所以任何路径的 teamfulness 都是 1。树中共有 4 条长度为 4 的有趣路径，因此答案为 4。

样例 3

Sample grader 输入

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

输出

11

最长路径长度为 3，共有 4 条有趣路径，其中 3 条 teamfulness 为 3，1 条为 2，因此答案为

$3\cdot3+2=11$。

子任务

子任务	分值	$N$	$K$	额外限制
0		-		样例
1	4	$\le10^6$	$\le N$	每个地点直接相连的地点数至多为 2
2	7			存在一个地点与其他所有地点直接相连
3	9	$\le200$		无
4	10	$\le2\cdot10^3$
5		$\le10^6$	$=1$
6	9		$\le2$
7	11	$\le2\cdot10^5$	$\le50$
8	12		$\le N$
9	13	$\le10^6$		有趣路径的长度为奇数
10	15			无

Sample grader

输入格式：

• 第一行两个整数 $N,K$；
• 第二行 $N$ 个整数 $a_0,a_1,\dots,a_{N-1}$；
• 接下来 $N-1$ 行，第 $i$ 行两个整数 $u_i,v_i$。

输出一行一个整数，即函数返回值。