P17275 [eJOI 2026] Teamfulness

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: ```cpp long long teamfulness(int N, int K, std::vector a, std::vector u, std::vector 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.

### 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