Anton estaba paseando recientemente por Kaunas durante la EJOI cuando descubrió $N$ lugares donde los participantes pasan el rato. Los lugares están numerados del 0 al $N-1$. Cada lugar está ocupado por miembros de exactamente un equipo, y los equipos están numerados del 0 al $K-1$. Ten en cuenta que los miembros de un mismo equipo pueden ocupar una cantidad arbitraria de lugares. Algunos equipos pueden no ocupar ningún lugar.
Los lugares están conectados por $N-1$ carreteras de doble sentido de modo que existe exactamente un camino simple entre cualquier par de lugares, por lo que forman un árbol. Recuerda que un camino simple entre dos lugares es una secuencia de lugares distintos en la que cada dos lugares consecutivos están conectados por una carretera. La longitud de un camino es el número de carreteras que usa, es decir, uno menos que el número de lugares que visita.
Anton quiere dar un paseo por un camino simple, visitando tantos lugares como sea posible. Un camino simple se llama interesante si su longitud es la máxima posible entre todos los caminos simples del árbol. El teamfulness de un camino es el número de equipos distintos que Anton encuentra a lo largo de él.
Tu tarea es encontrar la suma de los teamfulness de todos los caminos interesantes diferentes. Dos caminos interesantes se consideran iguales si y solo si visitan exactamente el mismo conjunto de lugares; en particular, recorrer un camino en la dirección opuesta no produce un camino diferente.
Detalles de implementación
Debes implementar la siguiente función:
long long teamfulness(int N, int K, std::vector<int> a,
std::vector<int> u, std::vector<int> v)
- $N$: el número de lugares;
- $K$: el número de equipos;
- $a$: un arreglo de $N$ enteros, donde $a_i$ es el equipo que ocupa el lugar $i$, para cada $0 \le i < N$;
- $u, v$: arreglos de $N-1$ enteros, donde $u_i$ y $v_i$ son los dos lugares conectados por la $i$-ésima carretera, para cada $0 \le i < N-1$.
Esta función se llama exactamente una vez por caso de prueba y debe devolver la suma de los teamfulness de todos los caminos interesantes.
Restricciones
- $3 \le N \le 10^6$
- $1 \le K < N$
- $0 \le a_i < K$ para cada $0 \le i < N$
- $0 \le u_i, v_i < N$ para cada $0 \le i < N-1$
Ejemplos
Entrada 1
6 3 1 0 0 1 2 1 0 1 0 2 0 3 0 4 0 5
Salida 1
21
Nota 1
La longitud máxima de un camino simple es 2 (2 carreteras, 3 lugares), por lo tanto los caminos interesantes tendrán longitud 2. Hay dos caminos interesantes con teamfulness 3, siete caminos interesantes con teamfulness 2 y un camino interesante con teamfulness 1, para una suma total de 21.
Entrada 2
7 1 0 0 0 0 0 0 0 0 1 0 2 1 3 1 4 2 5 2 6
Salida 2
4
Nota 2
Los lugares y las carreteras que los conectan se ilustran a continuación (el equipo 0 está coloreado de amarillo):
El teamfulness de cada camino es 1 porque solo hay un equipo (el equipo 0) que está presente en todos los lugares. Hay 4 caminos interesantes (de longitud 4), por lo tanto la suma de teamfulness de todos los caminos interesantes también es 4.
Entrada 3
6 3 0 1 2 0 1 2 0 1 1 2 2 3 1 4 2 5
Salida 3
11
Nota 3
Los lugares y las carreteras que los conectan se ilustran a continuación (el equipo 0 está coloreado de amarillo, el equipo 1 de verde y el equipo 2 de rojo):
Los caminos interesantes tienen longitud 3. Hay cuatro caminos interesantes en total: tres tienen teamfulness 3 y uno tiene teamfulness 2. Por lo tanto, la suma total de teamfulness de todos los caminos interesantes es 11.
Subtareas
| Subtarea | Puntos | $N$ | $K$ | Restricciones adicionales |
|---|---|---|---|---|
| 0 | 0 | - | - | Los ejemplos. |
| 1 | 4 | $\le 10^6$ | $| Cada lugar está conectado directamente con a lo sumo otros 2 lugares. |
|
| 2 | 7 | $\le 10^6$ | $| Hay un lugar que está conectado directamente con todos los demás lugares. |
|
| 3 | 9 | $< 200$ | $|
| |
| 4 | 10 | $< 2 \cdot 10^3$ | ||
| 5 | 10 | $\le 10^6$ | $=1$ | |
| 6 | 9 | $\le 10^6$ | $<2$ | |
| 7 | 11 | $< 50$ | $|
| |
| 8 | 12 | $< 2 \cdot 10^5$ | $|
| |
| 9 | 13 | $\le 10^6$ | $| La longitud de un camino interesante es impar. |
|
| 10 | 15 | $< 10^6$ | $|
| |
Ejemplos
El formato de entrada es el siguiente:
- línea 1: dos enteros: los valores de $N$ y $K$;
- línea 2: $N$ enteros $a_0, a_1, \dots, a_{N-1}$, donde $a_i$ es el equipo que ocupa el lugar $i$;
- línea $3+i$: dos enteros $u$ y $v$: los dos lugares conectados por la $i$-ésima carretera.
El formato de salida es el siguiente:
- línea 1: un entero: el valor de retorno de la llamada.