[백준 1949] 우수 마을 (Swift)

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

 

DP & DFS

https://sio2whocode.tistory.com/296 SNS 얼리어답터 문제랑 완전 똑같은 문제임

차이점은 그때는 최소 얼리어답터 수를 구하는 거였다면(이 문제로 치면 최소 우수마을의 수)

지금은 우수마을의 주민 수의 최대 합을 구하는 문제

 

접근방법

트리모양에서 DP를 적용해서 푼다고 생각하면, 나와 내 자식노드까지가 sub tree라고 보고

내가 우수마을일 때, 내 자식노드들은 우수마을이면 안됨

내가 우수마을이 아닐 때, 내가 자식 노드들은 우수마을이어도 되고, 아니어도 된다.

이게 주된 로직이다. 다만 여기서는 내 인접 마을 중 하나는 우수마을이어야 되는데,

내가 우수마을이 아닐 때, 내 자식 노드들이 우수마을이 아니어도 된다고해서 내 부모 노드도, 내 자식노드도 우수마을이 아닌 결과가 나오면 어쩌나. 하고 고민을 했다.

근데 트리를 그려서 시뮬레이션을 돌려보면 우수마을인 주민수의 합이 최대가 되도록 하는 로직이라, 두 구간 이상 우수마을이 아닐 수가 없다. 연달아 우수구간이 아닌 경우는 그중 하나가 우수마을인 경우에 비해 확실히 주민 수의 합이 작을 테니까.

따라서 위의 로직만 지키면서 dfs를 진행하면 dp의 루트노드에 최종 결과가 담기게 된다. 

 

소스코드

import Foundation

func solution() -> Int {
    var answer:Int = 0
    
    let n:Int = Int(readLine()!)!
    let pop:[Int] = [0] + (readLine()?.split(separator: " ").map { Int($0)! })!
    
    
    var adj:[[Int]] = [[Int]](repeating: [], count: n+1)
    
    for _ in 0..<n-1 {
        let edge:[Int] = readLine()?.split(separator: " ").map { Int($0) } as! [Int]
        adj[edge[0]].append(edge[1])
        adj[edge[1]].append(edge[0])
    }
    
    var visit:[Bool] = [Bool](repeating: false, count: n+1)
    var dp:[[Int]] = [[Int]](repeating: [Int](repeating: 0, count: 2), count: n+1)
    // dp [0] -> 우수마을, [1] -> 우수마을 아님
    
    func getMax(_ start: Int) {
        visit[start] = true
        
        dp[start][0] = pop[start]
        
        for ch in adj[start] {
            // start가 우수 마을일 때 -> 인접 마을은 우수마을이면 X
            // start가 우수 마을 아니면 -> 인접 마을은 우수마을 일 수도, 아닐수도
            if !visit[ch] {
                getMax(ch)
                dp[start][0] += dp[ch][1]
                dp[start][1] += max(dp[ch][0], dp[ch][1])
            }
        }
    }
    getMax(1)
        
    answer = max(dp[1][0], dp[1][1])
    return answer
}

print(solution())