2023/08/16(수요일)

백준 24230(g5) - 트리 색칠하기

문제

정점이 N개인 트리가 있다. 정점에는 1부터 N까지 번호가 붙어있다. 트리의 루트는 항상 1번 정점이며 맨 처음에는 모든 정점이 하얀색으로 칠해져 있는 상태이다.

하나의 정점에 색칠하면 해당 정점 아래 있는 모든 정점이 같은 색으로 칠해진다. 색은 섞이지 않고 색칠할 때마다 그 색으로 덮어진다. 단, 하얀색으로 색칠할 수는 없다.

아래 그림처럼 정점 10개로 구성된 트리가 있다고 가정을 해보자.

https://upload.acmicpc.net/d60753e6-25d7-4baa-94c3-e99d84ed3d42/-/preview/

[그림 1] 하얀색으로 칠해져 있는 트리

3번 정점을 노란색으로 칠하면 그 아래 있는 정점 5, 6, 8, 9, 10 모두 노란색으로 칠해진다.

https://upload.acmicpc.net/6345e2aa-2c8a-4f59-a274-e9073c61e520/-/preview/

[그림 2] 정점 3에 노란색을 칠한 후 트리의 상태

그리고 정점 5에 파란색을 칠한다면 그 아래 있는 정점 8, 9, 10 모두 파란색으로 칠해진다.

https://upload.acmicpc.net/25b335ab-1493-4ca6-a4a0-87486da7e39d/-/preview/

[그림 3] 정점 5에 파란색을 칠한 후 트리의 상태

입력으로 트리의 정보와 정점의 색 정보가 주어진다. 색 정보는 음이 아닌 정수로 주어지며 값이 0인 경우는 항상 하얀색을 의미한다.

하얀색을 제외한 색만 사용해서 모든 정점을 주어진 색으로 칠하고 싶을 때 최소 몇 번 색을 칠해야 모든 정점을 원하는 색으로 칠할 수 있는지 구해보자.

1. 내 첫 풀이

import java.io.BufferedReader
import java.io.InputStreamReader
import java.util.*

fun main(args: Array<String>){
    val br = BufferedReader(InputStreamReader(System.`in`))
    val n = br.readLine().toInt()
    var inputColorNode = br.readLine().split(" ").map{ it.toInt() }
    var graph = Array<MutableList<Int>>(n){ mutableListOf() }
    var colorNode = Array<Int>(inputColorNode.size){ 0 }
    var count = 0

    repeat(n - 1){
        var (node1, node2) = br.readLine().split(" ").map { it.toInt() }
        graph[node1 - 1].add(node2 - 1)
        graph[node2 - 1].add(node1 - 1)
    }

    for(i in inputColorNode.indices){
        if(inputColorNode[i] != 0 && inputColorNode[i] != colorNode[i]){
            bfs(graph, inputColorNode, colorNode, inputColorNode[i], i)
            count++
        }
    }

    print(count)
}

private fun bfs(graph: Array<MutableList<Int>>, inputColorNode: List<Int>, colorNode: Array<Int>, color: Int, index: Int){
    var q: Queue<Int> = LinkedList()
    val visited = HashSet<Int>()

    visited.add(index)
    colorNode[index] = color
    q.offer(index)

    while(q.isNotEmpty()){
        var preNode = q.poll()

        graph[preNode].forEach {
            if(it !in visited && it > preNode){
                q.offer(it)
                colorNode[it] = color
            }
        }
    }
}

→ 주어진 그래프에서 자식의 값은 항상 부모의 값보다 클 것이라고 생각하였고 생각나는데로 코드를 끄적였다,,

위 코드에서의 핵심은 colorNode라는 새로운 배열을 만들어 0으로 초기화를 해 주었으며 bfs를 돌 때마다 colorNode는 inputColorNode라는 주어진 색칠이 된 변수랑 똑같이 되어간다.

다시말해

//입력
10
0 1 0 2 1 0 2 2 2
1
4
5
5
2
6
5
8
7

위의 값들이 입력으로 주어졌을 때

colorNode는

[0, 0, 0, 0, 0, 0, 0, 0, 0, 0] → [0, 0, 1, 0, 1, 1, 0, 1, 1, 1] → [0, 0, 1, 0, 2, 1, 0 ,2, 2, 2] 로 변한다.

하지만 위의 방식은 시간초과를 생각하지 않은 방식으로 제출했을 때 시간초과가 났다….😓

2. 두 번째 풀이

import java.io.BufferedReader
import java.io.InputStreamReader
import java.util.*

fun main(args: Array<String>){
    val br = BufferedReader(InputStreamReader(System.`in`))
    val n = br.readLine().toInt()
    var colorNode = br.readLine().split(" ").map{ it.toInt() }
    var graph = Array<MutableList<Int>>(n){ mutableListOf() }
    var count = 0

    repeat(n - 1){
        var (node1, node2) = br.readLine().split(" ").map { it.toInt() }
        graph[node1 - 1].add(node2 - 1)
        graph[node2 - 1].add(node1 - 1)
    }

    if(colorNode[0] != 0){
        count++
    }

    print(bfs(graph, colorNode, count))
}

private fun bfs(graph: Array<MutableList<Int>>, colorNode: List<Int>, count: Int): Int{
    var q: Queue<Int> = LinkedList()
    val visited = HashSet<Int>()
    var cnt = count

    visited.add(0)
    q.offer(0)

    while(q.isNotEmpty()){
        var node = q.poll()

        graph[node].forEach {
            if(it !in visited){
                q.offer(it)
                visited.add(it)

                if(colorNode[it] != colorNode[node]){
                    cnt++
                }
            }
        }
    }

    return cnt
}

→ 두 번째 풀이의 핵심은 부모노드와 자식노드의 색깔 차이이다. 즉, 부모노드와 자식 노드의 색깔이 다르면 무조건 색을 칠해야만 한다는 뜻이고 색깔 차이가 났을 때마다 count를 하나씩 더해주면 된다.

여기서 유의할 점이 있는데, 위의 알고리즘은 첫 시작노드를 검사하지 않는다. 그렇기 때문에

if(colorNode[0] != 0){
        count++
    }

의 코드를 통하여 루트노드가 0이 아닐시 count를 하나 더하여 bfs를 실행해주면 된다.