[Swift] 프로그래머스 LV2. 전력망을 둘로 나누기

프로그래머스 LV2. 전력망을 둘로 나누기

 

접근 방식

 

트리에서 간선 하나를 끊으면 두 개의 서브트리로 나뉨. 모든 간선을 하나씩 끊어보며 두 쪽의 노드 수 차이가 가장 작은 경우를 찾는 문제.

 

input은 아래 코드처럼 양방향으로 받음. 그래야 각 정점의 연결 관계가 빠지지 않음.

 

이중 for문으로, 상위 for문은 정점을 순회하고 하위 for문은 그 정점에 연결된 간선을 돌며 removeFirst()로 하나씩 끊음. 이 문제는 트리를 전제하므로 간선 하나를 끊으면 끊긴 반대편은 다시 방문할 수 없게 됨. (그래프라면 반대쪽 정점에서도 끊어주는 작업이 필요함.)

 

끊은 뒤 그 정점에서 DFS로 방문 가능한 노드를 셈. 끊긴 쪽을 못 가므로 count는 전체 정점 수보다 작아지고, n - 2*count의 절댓값이 두 그룹의 크기 차이가 됨. 확인한 간선은 다시 append해 원상 복구함.

 

시행착오

 

문제가 되는 부분은 minValue를 업데이트하는 과정에서 n - count - count를 해야 하는데, 어째서인지 count - n과 동일한 값이라고 생각했음. 이를 수정하여 해결함.

 

코드

 

    func solution(_ n:Int, _ wires:[[Int]]) -> Int {
        var list = [Int :[Int]]()
        for i in 0..<wires.count {
            let start = wires[i][0]
            let end =  wires[i][1]
            
            if list[start] == nil {
                list[start] = [end]
            } else {
                list[start]?.append(end)
            }
            
            if list[end] == nil {
                list[end] = [start]
            } else {
                list[end]?.append(start)
            }
        }
        
        // 1. 전력망을 끊는다.
        // 2. DFS / BFS를 사용하여 방문한 지점을 카운팅한다.
        // 3. 최소값을 계산한다.
        var minValue = Int.max
        
        for i in 1...n {
            for _ in 1...list[i]!.count {
                let point = list[i]!.removeFirst() // 전력망 끊기
                
                if point != i {
                    let count = DFS(graph: list, start: i).count
                    minValue = min(abs(n - count - count), minValue)
                }
                list[i]!.append(point)
            }
        }
        
        
        return minValue
    }
    
func DFS<T> (graph: [T: [T]], start: T) -> [T] {
    var visitedQueue: [T] = []
    var needVisitStack: [T] = [start]
    
    while !needVisitStack.isEmpty {
        let node: T = needVisitStack.removeLast()
        if visitedQueue.contains(node) { continue }
        
        visitedQueue.append(node)
        needVisitStack += graph[node] ?? []
    }
    
    return visitedQueue
}