서현이가 풀지 못한 문제

View as PDF

Submit solution

Points: 100
Time limit: 1.0s
Memory limit: 256M

Authors:
Problem types

문제

ANAGETDON 2025 전날까지 킬러 문제를 만들지 못한 서현이는 결국, 자신이 풀지 못했던 문제를 출제하기로 했다.

서현이는 이 문제를 풀지 못했다. 여러분은 이 문제를 풀 수 있을까?


\(N\)개의 정점으로 이루어진 트리가 주어진다. 각 정점은 \(1\)부터 \(N\)까지 번호가 매겨져 있으며, 빨간색 또는 파란색 중 하나로 색칠되어 있다.

이때, 다음 조건을 만족하는 서로 다른 네 정점 쌍 \((A, B, C, D)\)의 경우를 세는 프로그램을 작성해보자.

  • \(A-B-C-D\)는 트리 상의 경로다.
  • 정점들의 색깔 순서는 다음 두 경우 중 하나와 일치해야 한다.
    • 빨간색 - 파란색 - 파란색 - 빨간색
    • 파란색 - 빨간색 - 빨간색 - 파란색

단, \((A, B, C, D)\)와 그의 역순인 \((D, C, B, A)\)는 같은 경우로 본다.

입력

첫째 줄에 정점의 개수 \(N(1\le N\le 50\,000)\)이 주어진다.

둘째 줄에는 \(i(1\le i\le N)\)번째 정점의 색깔이 공백으로 구분되어 주어진다. \(0\)이면 빨간색, \(1\)이면 파란색으로 색칠되어 있다는 의미다.

셋째 줄부터 \(N - 1\)개의 줄에 걸쳐 간선 정보가 \(u\ v(1\le u, v\le N; u \ne v)\)의 형태로 주어진다. \(u\)와 \(v\)를 잇는 양방향 간선이 존재한다는 의미다.

출력

조건을 만족하는 정점 네 쌍의 경우의 수를 출력한다.

예제 입력 1

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

예제 출력 1

1

가능한 경우는 \(1 - 2 - 4 - 5\) 단 하나다. \(5-4-2-1\)과 \(1 - 2 - 4 - 5\)는 같은 경우로 보기 때문에 \(1\)을 출력한다.


Comments

There are no comments at the moment.