[baekjoon] #1949: 우수 마을

problem link

• If i(th) village is excellent village, the adjacent village (child villages) must not be excellent village.
• If i(th) village is not excellent village, the adjacent village (child villages) can be excellent village and not be excellent village.
• So, to get the maximum resident of excellent village, use memoization.

lateinit var tree: Array<MutableList<Int>>
lateinit var dp: Array<Array<Int>>

fun main() {
    val br = System.`in`.bufferedReader()
    val bw = System.out.bufferedWriter()
    
    val n = br.readLine().toInt()
    val residents = br.readLine().split(" ").map { it.toInt() }
    tree = Array(n + 1) { mutableListOf() }
    for (i in 0 until n - 1) {
        val (v1, v2) = br.readLine().split(" ").map { it.toInt() }
        tree[v1].add(v2)
        tree[v2].add(v1)
    }

    // dp[i][0] : i(th) village is not excellent village
    // dp[i][1] : i(th) village is excellent village
    dp = Array(n + 1) { Array(2) { 0 } }
    for (i in 1..n) {
        dp[i][1] = residents[i - 1]
    }
    dfs(1, 0)
    bw.write("${maxOf(dp[1][0], dp[1][1])}")

    br.close()
    bw.close()
}

fun dfs(num: Int, parent: Int) {
    for (child in tree[num]) {
        // top-down approach
        if (child != parent) {
            dfs(child, num)
            // child can not be excellent village
            dp[num][1] += dp[child][0]
            // child might be excellent village or not
            dp[num][0] += maxOf(dp[child][0], dp[child][1])
        }
    }
}

Time Complexity

O(V + E) (V: n, E: n-1)
O(n)