트리 색칠하기(백준 24230)

https://www.acmicpc.net/problem/24230

24230번: 트리 색칠하기

 

처음에는 트리를 순회하며 색을 칠하려 했으나 시간 초과로 인해 다른 방법을 찾아야 했다. 다른 풀이법을 보니 부모와 자식 노드의 색이 다르면 색을 1회 칠한 것으로 처리해서 색칠 횟수를 구하면 된다. 추가로 트리의 가장 첫 번째 노드가 색이 흰색이 아니면 정답에 1을 더한다. 

import sys
# 표준 입력을 더 빠르게 받기 위해 sys.stdin.readline 사용
input = sys.stdin.readline

# 노드의 개수 입력 받기
num_nodes = int(input())
# 각 노드의 색상 정보 입력 받기
node_colors = list(map(int, input().rstrip().split()))
# 필요한 색칠 횟수 초기화
num_color_changes = 0

# 노드 간 연결 정보 입력 받기 (트리 구조이므로, 노드 개수 - 1만큼의 연결 정보가 존재)
for _ in range(num_nodes - 1):
    # 연결된 두 노드의 번호 입력 받기
    node_a, node_b = map(lambda x: int(x) - 1, input().rstrip().split())
    # 두 노드의 색상이 다르면 색칠 횟수 증가
    if node_colors[node_a] != node_colors[node_b]:
        num_color_changes += 1

# 첫 번째 노드(루트 노드)의 색상이 0이 아닌 경우, 색칠 횟수 추가 증가
if node_colors[0] != 0:
    num_color_changes += 1

# 필요한 색칠 횟수 출력
print(num_color_changes)