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

https://school.programmers.co.kr/learn/courses/30/lessons/86971?language=swift

 

DFS, 그래프

트리 모양의 그래프에서 간선을 하나 끊고 두 서브트리로 나눌 때,

이 두 서브트리의 노드 수의 차의 최솟값을 구하는 문제

 

접근방법

1. 양방향 그래프의 인접리스트를 생성 (서브트리 탐색 위함)

2. 모든 전선을 하나씩 끊어가면서, 각 경우에 대해 두 서브 트리의 노드 개수 차이 최솟값 갱신

2-1. 끊는 전선에 연결되어 있는 노드를 루트노드로 한 두 서브트리 중 하나의 서브트리 노드 개수 구하는 함수 선언 (getChildCnt - dfs)

 

- visit 배열을 inout으로 사용

- childCnt를 구하는 과정에서 초기값을 1로 해주어야 본인(서브트리의 루트노드)를 포함할 수 있음

 

시간 복잡도

• wires.count == n - 1 (트리 구조)
• 각 전선을 하나씩 제거 → O(n)
• 제거 후 DFS로 노드 수 계산 → O(n) (최악의 경우 전 노드 탐색)

전체 시간복잡도: O(n²) (n ≤ 100 → 충분히 가능)

 

소스코드

import Foundation

func solution(_ n:Int, _ wires:[[Int]]) -> Int {
    
    // 양방향 그래프 인접리스트 생성
    var adj:[[Int]] = [[Int]](repeating: [], count: n+1)
    for wire in wires {
        let a = wire[0]
        let b = wire[1]
        
        adj[a].append(b)
        adj[b].append(a)
    }
    
    // 서브트리 노드 수 구하는 재귀 함수
    func getChildCnt(_ parent:Int, _ visit: inout [Bool]) -> Int {
        var childCnt:Int = 1 // 본인
        
        for child in adj[parent] {
            if !visit[child] {
                visit[child] = true
                childCnt += getChildCnt(child, &visit)
            }
        }
        
        return childCnt
    }
    
    var result:Int = Int.max
    
    // 전선 하나씩 끊어가면서 서브트리 노드 수 차이 최솟값 갱신
    for wire in wires {
        
        var visit:[Bool] = Array<Bool>(repeating: false, count: n+1)
        
        visit[wire[0]] = true
        visit[wire[1]] = true
        var aTreeCnt:Int = getChildCnt(wire[0], &visit)
        var bTreeCnt:Int = n-aTreeCnt
        
        result = min(result, abs(aTreeCnt - bTreeCnt))
    }
    
    return result
}