[BOJ/Python, Java] 백준 24230번 - 트리 색칠하기

백준 #24230 트리 색칠하기

24230번: 트리 색칠하기

 

Python ver.

from collections import deque
def bfs():
    nowColor = [0] * n
    visited = [False] * n
    cnt = 0
    q = deque()
    q.append(0)
    visited[0] = True
    while q:
        w = q.popleft()
        if nowColor[w] != color[w] and color[w] != 0:
            nowColor[w] = color[w]
            cnt += 1
        for v in grid[w]:
            if visited[v]: continue
            nowColor[v] = nowColor[w]
            visited[v] = True
            q.append(v)
    return cnt
            

n = int(input())
color = list(map(int, input().split()))
grid = [[] for _ in range(n)]
for _ in range(n-1):
    v1, v2 = map(int, input().split())
    grid[v1-1].append(v2-1)
    grid[v2-1].append(v1-1)
print(bfs())

 

Java ver.

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.StringTokenizer;

public class Main {
	
	static int n;
	static int[] color;
	static ArrayList<ArrayList<Integer>> grid;
	
	public static int bfs() {
		int[] nowColor = new int[n];
		boolean[] visited = new boolean[n];
		ArrayDeque<Integer> q = new ArrayDeque<Integer>();
		int cnt = 0;
		q.add(0);
		visited[0] = true;
		
		while (!q.isEmpty()) {
			int w = q.pollFirst();
			if (nowColor[w] != color[w] && color[w] != 0) {
				nowColor[w] = color[w];
				cnt++;
			}
			for (int i = 0; i < grid.get(w).size(); i++) {
				int v = grid.get(w).get(i);
				if (visited[v]) {
					continue;
				}
				nowColor[v] = nowColor[w];
				visited[v] = true;
				q.add(v);
			}
		}
		return cnt;
		
	}

	public static void main(String[] args) throws IOException {
		
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		StringTokenizer st;
		
		// 입력 받기 
		n = Integer.parseInt(br.readLine());  // 정점의 개수
		color = new int[n];  // 각 색 정보
		grid = new ArrayList<ArrayList<Integer>>();
		for (int i = 0; i < n; i++) {
			grid.add(new ArrayList<Integer>());
		}
		st = new StringTokenizer(br.readLine());
		for (int i = 0; i < color.length; i++) {
			color[i] = Integer.parseInt(st.nextToken());
		}
		
		for (int i = 0; i < n-1; i++) {
			st = new StringTokenizer(br.readLine());
			int v1 = Integer.parseInt(st.nextToken());
			int v2 = Integer.parseInt(st.nextToken());
			grid.get(v1-1).add(v2-1);
			grid.get(v2-1).add(v1-1);
			
		}
		
		
		System.out.println(bfs());
		
		
		
		
		
	}  // end of main
}  // end of class

부모 노드를 칠하면 자식 노드도 자연스럽게 칠해지기 때문에 트리 상단부터 탐색하며 내려와야 한다. 

현재의 색칠 상태를 나타내는 nowColor와 원하는 색칠 상태를 저장하고 있는 color로, 부모 노드의 색칠 상태를 비교했을 때 원하는 색이 아니라면 원하는 색으로 칠해준다. (색칠 횟수 cnt + 1) 자식 노드는 queue에 담을 때 부모 노드와 같은 색으로 칠해준다.