[99클럽 코테 스터디 24일차 TIL] 전력망을 둘로 나누기 - 프로그래머스

트리 완전탐색이다.

DFS가 조금 더 맞는 키워드라고 생각한다.

 

https://school.programmers.co.kr/learn/courses/30/lessons/86971

n개의 송전탑이 전선을 통해 하나의 트리 형태로 연결되어 있습니다. 
당신은 이 전선들 중 하나를 끊어서 현재의 전력망 네트워크를 2개로 분할하려고 합니다. 
이때, 두 전력망이 갖게 되는 송전탑의 개수를 최대한 비슷하게 맞추고자 합니다.

송전탑의 개수 n, 그리고 전선 정보 wires가 매개변수로 주어집니다. 
전선들 중 하나를 끊어서 송전탑 개수가 가능한 비슷하도록 두 전력망으로 나누었을 때,
두 전력망이 가지고 있는 송전탑 개수의 차이(절대값)를 return 하도록 solution 함수를 완성해주세요.

[제한사항]
n은 2 이상 100 이하인 자연수입니다.
wires는 길이가 n-1인 정수형 2차원 배열입니다.
wires의 각 원소는 [v1, v2] 2개의 자연수로 이루어져 있으며, 
이는 전력망의 v1번 송전탑과 v2번 송전탑이 전선으로 연결되어 있다는 것을 의미합니다.
1 ≤ v1 < v2 ≤ n 입니다.
전력망 네트워크가 하나의 트리 형태가 아닌 경우는 입력으로 주어지지 않습니다.

출처: https://school.programmers.co.kr/learn/courses/30/lessons/86971

https://school.programmers.co.kr/learn/courses/30/lessons/86971

위 그림처럼 나누면 된다.

풀이 접근

먼저 연결된 노드를 파악하는 것이 중요하다고 생각했다.

var adjNode: [[Int]] = Array(repeating: [Int](), count: n + 1)
var visited: [Bool] = Array(repeating: false, count: n + 1)
for w in wires {
    adjNode[w[0]].append(w[1])
    adjNode[w[1]].append(w[0])
}

이후 각 노드마다 본인과 모든 자식 노드가 몇개의 노드로 구성되어있는지 DFS로 확인했다.

var nodeSize: [Int] = Array(repeating: 0, count: n + 1)
func DFSforChildCount(index: Int) -> Int {
    visited[index] = true
    var count = 1
    for i in adjNode[index] {
        if visited[i] { continue }
        count += DFSforChildCount(index: i)
    }
    nodeSize[index] = count
    return count
}
// 무조건 1부터 시작
nodeSize[1] = DFSforChildCount(index: 1)

 

본격적으로 하나씩 노드를 끊어보면서 개수를 비교한다.

var result = Int.max
visited = Array(repeating: false, count: n + 1)
func DFS(index: Int, parentSize: Int) {
    visited[index] = true
    for i in adjNode[index] {
        if visited[i] { continue }
        let cutSize = nodeSize[index] - nodeSize[i] + parentSize
        result = min(result, abs(cutSize - nodeSize[i]))
        DFS(index: i, parentSize: cutSize)
    }
}
DFS(index: 1, parentSize: 0)
return result

두개의 DFS코드로 해결되는 문제였다.

최종 코드 - Swift

import Foundation

func solution(_ n:Int, _ wires:[[Int]]) -> Int {
    var adjNode: [[Int]] = Array(repeating: [Int](), count: n + 1)
    var visited: [Bool] = Array(repeating: false, count: n + 1)
    var nodeSize: [Int] = Array(repeating: 0, count: n + 1)
    for w in wires {
        adjNode[w[0]].append(w[1])
        adjNode[w[1]].append(w[0])
    }
    
    func DFSforChildCount(index: Int) -> Int {
        visited[index] = true
        var count = 1
        for i in adjNode[index] {
            if visited[i] { continue }
            count += DFSforChildCount(index: i)
        }
        nodeSize[index] = count
        return count
    }
    nodeSize[1] = DFSforChildCount(index: 1)
    
    var result = Int.max
    visited = Array(repeating: false, count: n + 1)
    func DFS(index: Int, parentSize: Int) {
        visited[index] = true
        for i in adjNode[index] {
            if visited[i] { continue }
            let cutSize = nodeSize[index] - nodeSize[i] + parentSize
            result = min(result, abs(cutSize - nodeSize[i]))
            DFS(index: i, parentSize: cutSize)
        }
    }
    DFS(index: 1, parentSize: 0)
    return result
}

최종 코드 - Python3

def solution(n, wires):
    adjNode = [[] for _ in range(n+1)]
    visited = [False for _ in range(n+1)]
    nodeSize = [0 for _ in range(n+1)]
    for w in wires:
        adjNode[w[0]].append(w[1])
        adjNode[w[1]].append(w[0])
        
    def DFSforChildCount(visited, nodeSize, adjNode, index):
        visited[index] = True
        count = 1
        for i in adjNode[index]:
            if visited[i]: continue
            count += DFSforChildCount(visited, nodeSize, adjNode, i)
        nodeSize[index] = count
        return count
    
    nodeSize[1] = DFSforChildCount(visited, nodeSize, adjNode, 1)
    
    result = 100
    visited = [False for _ in range(n+1)]
    
    def DFS(visited, nodeSize, adjNode, index, parentSize):
        nonlocal result
        visited[index] = True
        for i in adjNode[index]:
            if visited[i]: continue
            cutSize = nodeSize[index] - nodeSize[i] + parentSize
            result = min(result, abs(cutSize - nodeSize[i]))
            DFS(visited, nodeSize, adjNode, i, cutSize)

    DFS(visited, nodeSize, adjNode, 1, 0)
    return result

Conclusion

항상 Swift로 중첩 함수를 정의하여 문제를 풀 때는 Python으로 변환하기 조금 어렵다.