Anton은 최근 EJOI 기간 동안 Kaunas를 걷다가 참가자들이 모이는 장소 $N$개를 발견했다. 장소에는 0부터 $N-1$까지 번호가 붙어 있다. 각 장소는 정확히 한 팀의 멤버들이 차지하고 있으며, 팀은 0부터 $K-1$까지 번호가 붙어 있다. 같은 팀의 멤버들은 임의로 많은 장소를 차지할 수 있다. 어떤 팀은 장소를 전혀 차지하지 않을 수도 있다.
장소들은 $N-1$개의 양방향 도로로 연결되어 있어서, 임의의 두 장소 사이에 단순 경로가 정확히 하나 존재한다. 즉, 이들은 트리를 이룬다. 두 장소 사이의 단순 경로란, 서로 다른 장소들의 수열로서, 수열에서 연속하는 두 장소는 모두 도로로 연결되어 있는 경로를 뜻한다. 경로의 길이는 그것이 사용하는 도로의 수, 즉 방문하는 장소의 수보다 1 작은 값이다.
Anton은 가능한 한 많은 장소를 방문하면서 단순 경로를 따라 걷고 싶어 한다. 단순 경로의 길이가 트리의 모든 단순 경로 중 가능한 최댓값일 때, 그 경로를 흥미로운 경로(interesting path)라고 부른다. 경로의 teamfulness는 Anton이 그 경로를 따라 만나는 서로 다른 팀의 수이다.
여러분의 임무는 서로 다른 모든 흥미로운 경로에 대한 teamfulness의 합을 구하는 것이다. 두 흥미로운 경로는 정확히 같은 장소 집합을 방문할 때에만 같은 경로로 간주된다. 특히, 경로를 반대 방향으로 지나가는 것은 다른 경로를 만들지 않는다.
구현 세부사항
다음 함수를 구현해야 한다:
long long teamfulness(int N, int K, std::vector<int> a,
std::vector<int> u, std::vector<int> v)- $N$: 장소의 수;
- $K$: 팀의 수;
- $a$: $N$개의 정수로 이루어진 배열로, 각 $0 \le i < N$에 대해 $a_i$는 장소 $i$를 차지하는 팀이다;
- $u, v$: 각각 $N-1$개의 정수로 이루어진 배열로, 각 $0 \le i < N-1$에 대해 $u_i$와 $v_i$는 $i$번째 도로로 연결된 두 장소이다.
이 함수는 각 테스트에서 정확히 한 번 호출되며, 모든 흥미로운 경로에 대한 teamfulness의 합을 반환해야 한다.
제한
- $3 \le N \le 10^6$
- $1 \le K < N$
- 각 $0 \le i < N$에 대해 $0 \le a_i < K$
- 각 $0 \le i < N-1$에 대해 $0 \le u_i, v_i < N$
예제
입력 1
6 3 1 0 0 1 2 1 0 1 0 2 0 3 0 4 0 5
출력 1
21
참고
단순 경로의 최대 길이는 2(도로 2개, 장소 3개)이므로, 흥미로운 경로의 길이는 2이다. teamfulness가 3인 흥미로운 경로는 2개, teamfulness가 2인 흥미로운 경로는 7개, teamfulness가 1인 흥미로운 경로는 1개 있으며, 따라서 합은 21이다.
입력 2
7 1 0 0 0 0 0 0 0 0 1 0 2 1 3 1 4 2 5 2 6
출력 2
4
참고
장소들과 장소들을 연결하는 도로는 다음과 같다 (팀 0은 노란색으로 표시되어 있다):
모든 장소에는 한 팀(팀 0)만 있으므로 모든 경로의 teamfulness는 1이다. 흥미로운 경로는 (길이가 4인) 4개 있으며, 따라서 모든 흥미로운 경로에 대한 teamfulness의 합도 4이다.
입력 3
6 3 0 1 2 0 1 2 0 1 1 2 2 3 1 4 2 5
출력 3
11
참고
장소들과 장소들을 연결하는 도로는 다음과 같다 (팀 0은 노란색, 팀 1은 초록색, 팀 2는 빨간색으로 표시되어 있다):
흥미로운 경로의 길이는 3이다. 흥미로운 경로는 총 4개이며, 그중 3개는 teamfulness가 3이고 1개는 teamfulness가 2이다. 따라서 모든 흥미로운 경로에 대한 teamfulness의 합은 11이다.
서브태스크
| 서브태스크 | 점수 | $N$ | $K$ | 추가 제약 조건 |
|---|---|---|---|---|
| 0 | 0 | - | - | 예제들. |
| 1 | 4 | $\le 10^6$ | $| 모든 장소는 최대 2개의 다른 장소와 직접 연결되어 있다. |
|
| 2 | 7 | $\le 10^6$ | $| 어떤 장소 하나가 다른 모든 장소와 직접 연결되어 있다. |
|
| 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$ | $| 흥미로운 경로의 길이가 홀수이다. |
|
| 10 | 15 | $< 10^6$ | $|
| |
입력
입력 형식은 다음과 같다:
- 1번째 줄: 두 정수 – $N$과 $K$의 값;
- 2번째 줄: $N$개의 정수 $a_0, a_1, \dots, a_{N-1}$ – $a_i$는 장소 $i$를 차지하는 팀;
- $3+i$번째 줄: 두 정수 $u$와 $v$ – $i$번째 도로로 연결된 두 장소.
출력
출력 형식은 다음과 같다:
- 1번째 줄: 하나의 정수 – 호출의 반환값.