누텔라 트리 (Easy)

• 몇가지 관찰들. 트리에서 연결된 점들로 구성된 서브그래프를 찾으면 역시 트리이다. 트리에서 두 점 사이의 경로는 유일하다.
• 검은색 노드 B와 빨간색 노드 R 간에 엣지가 있으면 B-R의 누텔라 경로를 만들수 있다. R과 연결된 모든 빨간색 노드들 R_i에 대해서 B로 시작해서 R_i로 끝나는 누텔라 경로를 유일하게 만들수 있다.
• 그러면 이제 문제는 간단해진다. 빨간색 노드들로 구성된 커넥티드 컴포넌트들을 모두 찾고, 크기를 계산해둔다. 검은색 노드와 빨간색 노드간에 엣지가 발견되면, 그 빨간색 노드가 포함된 커넥티드 컴포넌트의 크기만큼 누텔라 경로를 추가할수 있다.
• 커넥티드 컴포넌트들을 찾는 것은 탐색으로도 가능하고 Disjoint Set을 이용해도 된다. Disjoint Set이 코딩이 좀더 간단하다. 이 경우 총 시간복잡도는 O(n*α(n))이 된다.

"""Solution code for "BOJ 23040. 누텔라 트리 (Easy)".

- Problem link: https://www.acmicpc.net/problem/23040
- Solution link: http://www.teferi.net/ps/problems/boj/23040

Tags: [DisjointSet]
"""

import sys
from teflib import disjointset


def main():
    N = int(sys.stdin.readline())
    edges = [sys.stdin.readline() for _ in range(N - 1)]
    color = sys.stdin.readline()

    dsu = disjointset.DisjointSet(N)
    red_groups = []
    for edge in edges:
        u, v = [int(x) for x in edge.split()]
        if color[u - 1] == color[v - 1] == 'R':
            dsu.union(u - 1, v - 1)
        elif color[u - 1] == 'R':
            red_groups.append(u - 1)
        elif color[v - 1] == 'R':
            red_groups.append(v - 1)

    print(sum(dsu.size(x) for x in red_groups))


if __name__ == '__main__':
    main()

• Dependency: teflib.disjointset.DisjointSet