[프로그래머스] 전력망 둘로 나눈기(Python)

문제 링크

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

프로그래머스

 

문제 이해

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

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

 

제한사항

• n은 2 이상 100 이하인 자연수입니다.
• wires는 길이가 n-1인 정수형 2차원 배열입니다.

문제를 처음 이해했을 때는 어떻게 풀어야할 지 이해가 되지않았지만 BFS 방식을 사용하면 쉽게 풀 수 있을 것이라고 생각했다.

1. 그래프에 접합 노드들을 전부 추가한다.

2. BFS 함수를 만든 후 인접한 노드들의 갯수를 구하는 함수를 만든다.

3. wires에 있는 연결 노드들을 순회하며 하나씩 연결을 끊고(remove함수) bfs로 갯수의 차이를 구한다.

4. 연결을 끊은 노드들은 다시 연결시켜준다.

 

*이렇게 하면 시간 초과가 날 줄 알았는데 값이 잘 나왔다(?)

 

 

코드

from collections import deque
def solution(n, wires):
    answer = -1
    graph = [[] for _ in range(n+1)]
    for wire in wires:
        graph[wire[0]].append(wire[1])
        graph[wire[1]].append(wire[0])
        
    def bfs(val):
        visited = [0 for _ in range(n+1)]
        q = deque()
        q.append(val)
        cnt = 1
        visited[val] = 1
        
        while q:
            node = q.popleft()
            for tmp in graph[node]:
                if visited[tmp] == 0:
                    q.append(tmp)
                    visited[tmp] = 1
                    cnt += 1
        return cnt
    res = n
    for wire in wires:
        graph[wire[0]].remove(wire[1])
        graph[wire[1]].remove(wire[0])
        
        res = min(abs(bfs(wire[0])-bfs(wire[1])),res)
        
        graph[wire[0]].append(wire[1])
        graph[wire[1]].append(wire[0])
        
    return res