[백준/BOJ] - 24230번 python 풀이 - dfs

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

24230번: 트리 색칠하기

---

---

색을 칠하면 자식 노드들이 모두 색칠해진다.

0은 디폴트 값이다. (하얀색)

 

노드들 사이의 관계 매핑시에 양방향 매핑 필요.

 

dfs 탐색을 위한 반복문 부분에서 0이 아닌 부분(하얀색이 아니고)이고 방문 처리 되지 않았을 경우

dfs 탐색 시작. - color list는 [0]에 대한 빈값 처리를 해주지 않았기에 color[i-1]

 

dfs 함수에선 연결된 노드들을 탐색 시도하고 연결된 노드들의 방문 여부와 색상 일치 여부 확인

dfs 탐색 1 사이클이 종료됐을 경우 cnt값 += 1

 

pypy3로 돌렸더니 해결완료.

---

n = int(input())
color = list(map(int,input().split()))
arr = [[] for _ in range(n+1)]
ok = [False] * (n+1)
cnt = 0
for _ in range(n-1):
    a,b = map(int,(input().split()))
    arr[a].append(b) #트리 연결 관계 매핑
    arr[b].append(a) #트리 연결 관계 매핑
def dfs(n,c):    
    ok[n] = True
    for i in arr[n]:
        if not ok[i] and color[i-1] == c: #방문하지 않았고 색상이 일치할 경우
            dfs(i,c) #탐색 시작

for i in range(1,n+1): #1~7
    if not ok[i] and color[i-1] != 0: #흰색이 아니고 방문하지 않았을 경우
        dfs(i,color[i-1]) #color은 [0]에 대한 값을 설정해주지 않아서 i-1
        cnt+=1 #탐색한번 끝내면 값 += 1 
print(cnt)