99클럽 - 전력망을 둘로 나누기

프로그래머스

아이디어

이 전선들 중 하나를 끊어서 현재의 전력망 네트워크를 2개로 분할하려고 합니다. 이때, 두 전력망이 갖게 되는 송전탑의 개수를 최대한 비슷하게 맞추고자 합니다.

위 조건을 보고 빠르게 떠올릴 수 있는 것은 다음과 같다.

• 트리를 만들어야 한다.
• 매번 탐색하는 건 문제가 좀 있다.
• 트리가 있다면 각 노드의 서브트리를 계산해 두는 건 많은 경우에서 이득이다.
• 그리고 이 문제에서 서브트리를 계산해 두는 건 핵심이자 굉장히 많이 이득이다.

트리는 흔하지만 중요한 자료구조로 생성부터 탐색까지 자세히 정리해 본다.

풀이

tree = {}
for u, v in wires:
    if u not in tree:
        tree[u] = []
    if v not in tree:
        tree[v] = []
    tree[u].append(v)
    tree[v].append(u)

트리는 흔히 Dictionary로 만든다.
key에는 각각의 노드가 오고, value에는 간선(edge)로 연결 된 다른 노드를 저장해 연결 관계를 정의한다.
예를 들어 문제와 같은 배열이 들어가면 트리의 결과는 다음과 같다.

wires = [[1,3],[2,3],[3,4],[4,5],[4,6],[4,7],[7,8],[7,9]]

tree = {
    1: [3],
    2: [3],
    3: [1, 2, 4],
    4: [3, 5, 6, 7],
    5: [4],
    6: [4],
    7: [4, 8, 9],
    8: [7],
    9: [7]
}

이제 서브트리를 계산하기 위해 탐색을 진행해야 하는데 방식은 DFS와 BFS 두 가지가 있다.
BFS는 queue를 사용해 구현하므로 이번에도 import를 피하기 위해 DFS를 사용해 구현한다.
그리고, 서브트리는 '깊이'에 관련된 만큼 DFS 가 더 적합할 것이라는 생각이 든다.

subtree_sizes = {}

def dfs(node, parent):
    subtree_size = 1
    for neighbor in tree[node]:
        if neighbor == parent:
            continue
        size = dfs(neighbor, node)
        subtree_size += size
    subtree_sizes[node] = subtree_size
    return subtree_size

subtree_sizes에는 각 노드에 대한 서브트리의 갯수를 저장하게 된다.
dfs는 tree dictionry의 key를 순회하며 각각의 노드에 대한 계산을 진행한다.
파라미터로 호출된 부모 노드를 함께 전달해 이중 탐색을 막아주고, 더이상 탐색을 진행할 수 없다면 이 때 부터 1씩 반환해 최종적으로 subtree_sizes에 저장한다.
이에 따른 연산 과정은 다음과 같다.

• 노드 1 방문 시작:
  subtree_sizes는 비어 있다.
• 노드 3 방문 시작:
  탐색이 완료된 노드가 없어 subtree_sizes는 비어 있다.
• 노드 2 방문 및 완료:
  노드 2에서 더 이상 방문할 노드가 없으므로 자기 자신만 포함하여 크기가 1이 된다.
  subtree_sizes는 노드 2의 정보를 가지고 있다.
• 노드 4 방문 시작:
  노드 4에서 노드 5와 6, 7를 방문한다.
• 노드 5 방문 및 완료:
  노드 5 는 더 이상 방문할 노드가 없으므로 크기가 1아다.
  subtree_sizes는 노드 2, 5의 정보를 가지고 있다.
• 노드 6 방문 시작:
  노드 6도 마찬가지로 크기가 1이다.
  subtree_sizes는 노드 2, 5, 6의 정보를 가지고 있다.
• 노드 7 방문 시작:
  노드 7에서 노드 8과 9를 방문한다.
• 노드 8 방문 시작:
  노드 8은 자기 자신만 포함하여 크기가 1이다.
  subtree_sizes는 노드 2, 5, 6, 8의 정보를 가지고 있다.
• 노드 9 방문 시작:
  노드 9도 자기 자신만 포함하여 크기가 1이다.
  subtree_sizes는 노드 2, 5, 6, 8, 9의 정보를 가지고 있다.
• 노드 7 계산 완료:
  노드 7은 자신과 노드 8, 노드 9를 포함하여 크기가 3(1 + 1 + 1)이다.
  subtree_sizes는 노드 2, 5, 6, 7, 8, 9의 정보를 가지고 있다.
• 노드 4 계산 완료:
  노드 4는 자신, 노드 5, 노드 6, 노드 7(노드 8과 9 포함)을 포함하여 크기가 6(1 + 1 + 1 + 3)이다.
  subtree_sizes는 노드 2, 4, 5, 6, 7, 8, 9의 정보를 가지고 있다.
• 노드 3 계산 완료:
  노드 3은 자신, 노드 2, 노드 4(노드 5, 6, 7, 8, 9 포함)을 포함하여 크기가 8(1 + 1 +6)입니다.
  subtree_sizes는 노드 2, 3, 4, 5, 6, 7, 8, 9의 정보를 가지고 있다.
• 노드 1 계산 완료:
  노드 1은 전체 트리를 포함하여 크기가 9입니다.
  subtree_sizes는 노드 1, 2, 3, 4, 5, 6, 7, 8, 9의 정보를 가지고 있다.

과정이 길지만 결과는 단순하다.

subtree_sizes = {
    1: 9,
    2: 1,
    3: 8,
    4: 6,
    5: 1,
    6: 1,
    7: 3,
    8: 1,
    9: 1
}

이제는 간선의 정보에 해당하는 wires 배열을 순회하며 양쪽의 서브트리의 노드 수의 차가 가장 적은 경우를 골라 내면 된다.

min_diff = float('inf')
for u, v in wires:
    if subtree_sizes[u] > subtree_sizes[v]:
        size_v_side = subtree_sizes[v]
        size_u_side = n - size_v_side
    else:
        size_u_side = subtree_sizes[u]
        size_v_side = n - size_u_side

    diff = abs(size_u_side - size_v_side)
    min_diff = min(min_diff, diff)

노드들을 끊으면서 차이를 diff에 저장하고, min_diff와 비교하며 갱신해 최솟값을 찾는다.

def solution(n, wires):
    tree = {}
    for u, v in wires:
        if u not in tree:
            tree[u] = []
        if v not in tree:
            tree[v] = []
        tree[u].append(v)
        tree[v].append(u)
    
    subtree_sizes = {}
    
    def dfs(node, parent):
        subtree_size = 1
        for neighbor in tree[node]:
            if neighbor == parent:
                continue
            size = dfs(neighbor, node)
            subtree_size += size
        subtree_sizes[node] = subtree_size
        return subtree_size
    
    dfs(1, None)
    
    min_diff = float('inf')
    for u, v in wires:
        if subtree_sizes[u] > subtree_sizes[v]:
            size_v_side = subtree_sizes[v]
            size_u_side = n - size_v_side
        else:
            size_u_side = subtree_sizes[u]
            size_v_side = n - size_u_side
        
        diff = abs(size_u_side - size_v_side)
        min_diff = min(min_diff, diff)
    
    return min_diff

최종 코드는 위와 같다.