[Python] 24230번 트리 색칠하기

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

24230번: 트리 색칠하기

---

 

23/05/29

 

 

골드 5 짜리 문제 치고는 아이디어가 매우 좋았던 문제이다.

---

 

문제 접근 방식:

 

 

문제 정보를 보면, 트리가 주어져 있는데 각 노드별로 색깔 정보가 주어진다고 했다.

 

한 노드를 색칠하면, 그 노드를 포함한 서브트리는 모두 같은 색으로 색칠된다고 한다.

 

이때, 모든 정점을 원하는 색으로 칠하고자 할 때 구하는 최소 색칠 횟수를 구하는 것이 우리의 목적이다.

 

 

 

맨 처음 떠올렸던 방식은 BFS로 루트에서 시작하여 맨 위에서부터 탐색을 하는데, 그 만약 색깔이 칠해져있지 않은 상황인데 그 노드가 색깔이 칠해져 있다거나, 색칠되야 하는 색깔과 현재 색깔이 다르다면 그 노드를 포함한 서브트리를 직접 또다시 BFS로 색칠해 주는 방식이었다.

 

즉, 다시 말하자면 BFS안에 BFS를 돌리는 형태로, 문제 정보를 그대로 구현한 방식이라고 할 수 있다.

 

근데 이런 식으로 문제를 접근하면 시간 초과로 문제를 틀릴 수 있다.

 

BFS는 $O(V + E)$만큼의 시간 복잡도가 소요되는데, BFS안에 BFS를 돌린다면 $O((V+E)^2)$으로, 약 $400000^2$이 되어서 시간 내에 돌아가기 힘들기 때문이다.

 

 

 

때문에 탐색을 한 번만에 하여 구하는 방법을 고안해야 됐었다.

 

DFS를 이용한 방식을 떠올렸고, DFS를 할 때 한 노드와 다른 노드가 연결되어 있을 때 두 노드의 색깔이 다르다면 그 노드를 새롭게 색칠해줘야 한다는 아이디어를 떠올려서 실제로 그렇게 될 때마다 색칠 횟수를 $+1$을 해주었다.

 

또한, 처음 색칠하는 횟수는 루트 노드가 색칠돼있는 상태라면 $1$, 아니라면 $0$으로 놓고 시작하였다.

 

그렇게 한다면 탐색 한 번만에 색칠 횟수를 구할 수 있다는 결론을 얻을 수 있었다.

 

 

 

하지만 굳이 탐색을 하지 않아도 됐었다.

 

핵심 아이디어는 DFS의 풀이와 같다. 근데, 굳이 탐색을 하지 않아도 되는 이유는 간선들을 입력받을 때마다 저걸 체크해 주면 끝이기 때문이다.

 

 

 

요약하자면 다음과 같다.

 

어차피 색칠은 한번 하면 서브 트리의 모든 노드들이 다 색칠된다는 사실을 알고 있으므로, 굳이 새롭게 색칠하는 배열을 만들어서 색칠된 배열하고 비교하는, 즉, BFS안에 BFS를 돌리는 풀이는 하지 않아도 된다.

 

그러기 때문에 부모와 자식 간의 색깔이 같다면 그 자식은 부모 때문에 색칠된 것이라고 볼 수 있으므로 색칠 횟수를 세지 않는 것이다.

 

만약 다르다면 그 자식이 새롭게 색칠된 것이므로 색칠된 횟수를 세는 것이다.

 

이를 DFS를 통해 구할 수 있고, 사실은 굳이 탐색이 아니어도 이 그래프가 트리 하나로 이루어져 있다는 사실 때문에 그냥 그래프 간선 입력받을 때 세어주면 되는 것이다.

---

아래는 내가 위의 접근 방식과 같이 작성한 파이썬 코드이다. 더 보기를 누르면 확인할 수 있다.

# 탐색 풀이
# 24230번 트리 색칠하기
# dfs, 그래프 이론, 그래프 탐색, 트리
'''
접근 방법:
dfs를 돌린다.
돌릴 때, 색깔이 달라질 때마다 +1을 해준다.
'''
import sys
input = sys.stdin.readline

N = int(input())
color = list(map(int, input().rstrip().split()))
tree = [[] for _ in range(N)]
visited = [0 for _ in range(N)]

for _ in range(N-1):
    a, b = map(lambda x: int(x)-1, input().rstrip().split())
    tree[a].append(b)
    tree[b].append(a)

coloring_times = 0
stack = []
stack.append(0)
visited[0] = 1

while stack:
    node = stack.pop()
    for near_node in tree[node]:
        if visited[near_node] == 0:
            visited[near_node] = 1
            if color[node] != color[near_node]:
                coloring_times += 1
            stack.append(near_node)

if color[0] != 0:
    coloring_times += 1

print(coloring_times)

# 24230번 트리 색칠하기
# 트리
'''
접근 방법:
dfs돌릴 필요도 없이 연결된 간선끼리 색깔이 다르면 +1
'''
import sys
input = sys.stdin.readline

N = int(input())
color = list(map(int, input().rstrip().split()))
coloring_times = 0

for _ in range(N-1):
    a, b = map(lambda x: int(x)-1, input().rstrip().split())
    if color[a] != color[b]:
        coloring_times += 1

if color[0] != 0:
    coloring_times += 1

print(coloring_times)