프로그래머스 - 전력망을 둘로 나누기 [파이썬]

https://school.programmers.co.kr/learn/courses/30/lessons/86971?language=python3 

프로그래머스

---

풀이 과정

트리 형태의 전력망이 주어지는데, 전선을 하나 끊어서 두개의 전력망을 만들 때, 두 전력망의 송전탑의 차가 가장 적을 때의 차를 반환하는 문제이다.

 

위와 같은 전력망은 3-4 or 4-7을 끊으면 각 전력망내 송전탑의 차가 3으로 최소가 된다.

 

전선을 하나 끊고, 전선의 출발점과 도착점에 대해 BFS로 출발점이 속해 있는 망과 도착점이 속해 있는 망의 송전탑의 차를 모두 구해보면서, 최소인 경우를 반환하도록 하였다.

---

소스 코드

import collections

def solution(n, wires):
    def bfs(n, start, graph):
        visited = [False for _ in range(n+1)]
        visited[start] = True
        Q = collections.deque([start])
        network = 1
        
        while Q:
            cur = Q.popleft()
            for edge in graph[cur]:
                if not visited[edge]:
                    Q.append(edge)
                    visited[edge] = True
                    network += 1
                    
        return network
    
    graph = collections.defaultdict(set)
    for i, j in wires:
        graph[i].add(j)
        graph[j].add(i)
        
    answer = 101
    
    for vertex in graph:
        for edge in graph[vertex]:
            graph[vertex].remove(edge)
            graph[edge].remove(vertex)
            temp_vertex = bfs(n, vertex, graph)
            temp_edge = bfs(n, edge, graph)
            temp_answer = abs(temp_vertex - temp_edge)
            answer = min(answer, temp_answer)
            graph[vertex].add(edge)
            graph[edge].add(vertex)
    
    return answer