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

• 위클리 챌린지
• 전력망을 둘로 나누기
• Level 2

 

문제

• n개의 송전탑이 전선을 통해 하나의 트리 형태로 연결되어 있음
• 전선들 중 하나를 끊어 전력망 네트워크를 2개로 나누려 함
• 이때, 두 전력망이 갖게 되는 송전탑의 개수를 최대한 비슷하게 맞추려고 함
• 두 전력망이 갖는 송전탑 개수의 차이(절대값) 반환하기

 

• n은 2 이상 100 이하
• wires는 길이가 n-1 인 2차원 배열
• wires 원소는 [v1, v2] 2개의 자연수이며 각 수는 송전탑 번호를 의미
• 1 <= v1 <= v2 <= n

 

Idea

• 트리 정보를 저장
• 각 wire 의 두 노드에 각각 연결되어 있는 노드들의 개수를 구하고
• 연결된 노드 개수의 차이를 구한 다음
• 그 최솟값을 찾기

 

Code

Swift

• wire 를 끊을 때 이미 방문한 것 처럼 처리하면 연결된 노드 방향은 탐색하지 않을 수 있다.
• 또한 두 노드 중 하나에 연결된 노드들의 개수만 구하면, 나머지 노드에 연결된 것은 전체 개수에서 앞선 노드를 통해 구한 개수를 빼면 얻을 수 있다.
• 따라서 dfs 를 사용하여 wire 에 연결된 두 노드 중 한 노드에 연결된 노드들의 개수를 구하고
• 두 노드에 연결된 노드들의 개수 차이를 구하면 된다.

import Foundation

func solution(_ n:Int, _ wires:[[Int]]) -> Int {
    
    var tree = Array(repeating: [Int](), count: n+1)
    for edge in wires {
        tree[edge[0]].append(edge[1])
        tree[edge[1]].append(edge[0])
    }
    
    var visited = Array(repeating: false, count: n+1)
    func dfs(_ s: Int) -> Int {
        visited[s] = true
        var count = 0
        for next in tree[s] {
            if !visited[next] {
                count += dfs(next)
            }
        }
        visited[s] = false
        return count + 1
    }
    
    var answer = n
    for wire in wires {
        visited[wire[1]] = true
        let v1_count = dfs(wire[0])
        let v2_count = n - v1_count
        answer = min(answer, abs(v1_count - v2_count))
        visited[wire[1]] = false
    }
    
    return answer
}