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

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 입니다.
  • 전력망 네트워크가 하나의 트리 형태가 아닌 경우는 입력으로 주어지지 않습니다.

 

입출력 예

n	wires	result
9	[[1,3],[2,3],[3,4],[4,5],[4,6],[4,7],[7,8],[7,9]]	3
4	[[1,2],[2,3],[3,4]]	0
7	[[1,2],[2,7],[3,7],[3,4],[4,5],[6,7]]	1

 

 

입출력 예 설명

입출력 예 #1

• 다음 그림은 주어진 입력을 해결하는 방법 중 하나를 나타낸 것입니다.

• 4번과 7번을 연결하는 전선을 끊으면 두 전력망은 각 6개와 3개의 송전탑을 가지며, 이보다 더 비슷한 개수로 전력망을 나눌 수 없습니다.
• 또 다른 방법으로는 3번과 4번을 연결하는 전선을 끊어도 최선의 정답을 도출할 수 있습니다.

 

입출력 예 #3

• 다음 그림은 주어진 입력을 해결하는 방법을 나타낸 것입니다.

• 3번과 7번을 연결하는 전선을 끊으면 두 전력망이 각각 4개와 3개의 송전탑을 가지게 되며, 이 방법이 최선입니다.

 

문제 풀이

먼저 생각났던 풀이는 입력받은 전선의 정보를 통해 인접 리스트를 만든 뒤에, 전선의 정보를 하나씩 끊어보면서 각 송전탑의 개수를 세주는 방식이다. 이 방식으로 푼 뒤에 스터디를 진행하면서 개수를 세기 위한 더 좋은 방법들이 있다는 걸 알았고, 그 방식으로 수정하면서 점점 코드가 개선되었다. 

 

처음 나의 풀이

 1. 간선의 정보를 통해 인접 리스트를 만든다. 

 2. 간선을 하나씩 끊는다. (실제 인접 리스트에서 remove로 삭제 시켜준다.)

 3. 간선 하나를 끊은 인접 리스트에서 노드를 하나씩 돌아본다.

     방문하지 않는 노드를 시작점으로 주면 DFS로 해당 노드에 연결된 노드의 개수를 세준다.

 4. 3번에서 계산한 노드의 개수들을 리스트에 담아준다.(시작하는 노드가 다르면, 노드의 개수가 달라지므로)

 5. 개수의 차이 구하기

    :문제 설명에서 전력망 네트워크를 2개로 분할한다고 했으므로, 노드의 개수를 담은 리스트의 크기를 2로 예상했다. => 리스트에 담긴 2개의 원소의 차이를 구한다.

 6. 3번에서 반복문을 돌면서 3-5를 반복하며 원소의 차이중 최소값을 구한다.

 7. 2번에서 끊었던 간선을 다시 원상복구한다. (리스트에 append)

 8. 2-7번 과정을 반복한다.

from collections import deque

# 송전탑의 개수 세기
def count(x, graph, visited):
    cnt = 1
    queue = deque([x])
    visited[x] = True
    while queue:
        now = queue.popleft()
        for next in graph[now]:
            if not visited[next]:
                visited[next] = True
                queue.append(next)
                cnt += 1
    return cnt

def solution(n, wires):
    min_gap = int(1e9)
    graph = [[] for _ in range(n+1)]
    for w in wires:
        graph[w[0]].append(w[1])
        graph[w[1]].append(w[0])
    
    # 전선을 하나씩 끊어보며 완탐
    for w in wires:
        graph[w[0]].remove(w[1])
        graph[w[1]].remove(w[0])
        
        visited = [False] * (n+1)
        res = []
        # 각 전력망의 송전탑 개수 세기
        for i in range(1, n+1):
            if not visited[i]:
                res.append(count(i, graph, visited))
                
        # 차이의 최소값 구하기
        min_gap = min(min_gap, abs(res[0]-res[1]))
        
        # 원상복구
        graph[w[0]].append(w[1])
        graph[w[1]].append(w[0])
    
    return min_gap

 

 

이후 코드 리뷰를 진행하면서 이미 네트워크가 2개로 분리되는 것이 자명하니까, 굳이 간선을 끊은 인접 리스트의 모든 노드들을 확인하면서 2개 전력망의 노드 개수를 각각 구할 필요가 없다는 것을 알게 되었다.

그래서 아무 노드에서부터 시작해 해당 노드와 연결된 전력망의 크기를 구하면, 자동적으로 (전체 크기 - 해당 전력망의 크기) 가 나머지 전력망의 크기가 되는 것이다.

그래서 위 나의 코드 중 3-4번의 과정을 수정했다.

 

수정한 코드1 : 한쪽의 개수만 세기

 1. 간선의 정보를 통해 인접 리스트를 만든다. 

 2. 간선을 하나씩 끊는다. (실제 인접 리스트에서 remove로 삭제 시켜준다.)

 3. 1번 노드를 시작점으로 하여 BFS로 연결된 노드의 개수를 구한다. (여기서 개수를 담는 변수는 전역변수 설정)

 4. 해당 개수로 나머지 전력망의 개수도 구한 뒤, 두 전력망의 노드 개수의 차이값을 구한뒤, 최솟값과 비교해 갱신한다.

 5.  2번에서 끊었던 간선을 다시 원상복구한다. (리스트에 append)

 6. 2-5번 과정을 반복한다.

from collections import deque

# 송전탑의 개수 세기
def bfs(x, graph, visited):
    global cnt
    queue = deque([x])
    visited[x] = True
    while queue:
        now = queue.popleft()
        for next in graph[now]:
            if visited[next] == False:
                visited[next] = True
                cnt += 1
                queue.append(next)

def solution(n, wires):
    global cnt
    min_gap = int(1e9)
    graph = [[] for _ in range(n+1)]
    for w in wires:
        graph[w[0]].append(w[1])
        graph[w[1]].append(w[0])
    
    # 전선을 하나씩 끊어보며 완탐
    for w in wires:
        graph[w[0]].remove(w[1])
        graph[w[1]].remove(w[0])
        
        # 각 전력망의 송전탑 개수 세기
        visited = [False] * (n+1)
        cnt = 1
        bfs(1, graph, visited)
        # 차이의 최소값 구하기
        min_gap = min(min_gap, abs(visited.count(0)-1 - visited.count(1)))
        
        # 원상복구
        graph[w[0]].append(w[1])
        graph[w[1]].append(w[0])
    
    return min_gap

 

 

그러다 다른 팀원분이 인접 리스트에서 간선을 실제로 끊는 작업을 하지 않아도,

트리에서 루트를 바꿔주면서 개수를 세는 방식으로 같은 결과를 낼 수 있다고 하셨다.

처음에는 이 원리를 이해하지 못했는데, 계속 생각해보며 점차 이해를 할 수 있었다.

 

수정한 코드 2: 트리의 루트를 바꾸기

 1. 간선의 정보를 통해 인접 리스트를 만든다. 

 2. 각 노드를 돌면서 루트를 바꿔준다. (즉, 시작점을 1부터 n까지)

 3. DFS함수에서 자신을 포함해 자식 노드의 개수를 센다.

 4. 해당 개수를 리스트에 저장해준다.

 5. 개수 리스트를 돌면서 차이의 최소값을 구한다.

sum = [0] * (101)
visited = [False] * (101)
def dfs(num, graph):
    if sum[num] != 0:
        return sum[num]
    visited[num] = True
    sum[num] = 1
    for i in range(len(graph[num])):
        if not visited[graph[num][i]]:
            sum[num] += dfs(graph[num][i], graph)
    return sum[num]

def solution(n, wires):
    graph = [[] for _ in range(n+1)]
    for w in wires:
        graph[w[0]].append(w[1])
        graph[w[1]].append(w[0])
        
    min_gap = int(1e9)
    for i in range(1, n+1):
        cnt = dfs(i, graph)
        min_gap = min(abs((n-cnt)-cnt), min_gap)
        
    return min_gap

 

 

실행 시간 비교

처음 풀이	수정한 코드 1: 한 쪽의 개수만 세기

수정한 코드 2: 트리의 루트 바꾸기

 

확실히 전선을 직접 끊고 다시 원상복구하는 과정(remove, append)을 없애고

트리를 루트만 바꿔서 접근하며 노드의 개수를 구하는 알고리즘이 훨씬 빨랐다.

 

알고리즘 분류

• 완전탐색
• DFS
• BFS
• 트리