Python - [프로그래머스]86971-전력망을 둘로 나누기

출처

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

문제

연결된 하나의 트리가 주어지고, 임의의 한 간선을 끊을 때 양쪽의 송전탑 개수 차이가 가장 작은 경우를 찾는 문제이다.

접근 방법

완전 탐색 문제이다. 간선을 하나씩 전 끊어보고 가장 작은 경우를 찾아내야하는 문제이다. 간선을 하나씩 끊어서 현재 트리의 상태를 나타내는 배열을 만들고 끊어낸 두 송전탑을 기준으로 Deque를 활용한 BFS를 통해 양쪽 송전탑의 갯수를 센다.

풀이

하나씩 끊어내주는 메인함수와 BFS 두개의 함수로 구현하였다. 크게 어려운 부분이 없으니 코드를 바로 올린다.

from collections import deque

def bfs(flag, tree):
    queue = deque()
    queue.append(flag)
    visited = []
    answer = 0
    while queue:
        x = queue.popleft()
        visited.append(x)
        for i in range(len(tree[x])):
            if tree[x][i] != 0 and i not in visited:
                queue.append(i)
                answer += 1
    return answer


def solution(n, wires):
    answer = 101
    delete = 0
    for i in range(len(wires)):
        tree = [[0]*(n+1) for _ in range(n+1)]
        for j in range(len(wires)):
            if j != delete:
                tree[wires[j][0]][wires[j][1]] = 1
                tree[wires[j][1]][wires[j][0]] = 1
        a = bfs(wires[delete][0], tree)
        b = bfs(wires[delete][1], tree)
        delete += 1
        answer = min(answer, (abs(a - b)))
                
    
    return answer