[백준] 17471번: 게리맨더링 (파이썬/Python) - bfs, brute force

https://www.acmicpc.net/problem/17471

17471번: 게리맨더링

0) 그래프 입력 받기 : 파이썬의 defaultdict이라는 훌륭한 내장 라이브러리가 있는데, 이걸 활용하면 쉽다. 
그리고 문제에서 1부터 카운트하므로 -1 처리를 해주어서, 0부터 노드가 시작하게끔 처리해주었다. 이래야 나중에 index값으로 접근하기 용이하므로..

graph = defaultdict(list)

for i in range(n):
    data = list(map(int, input().split()))
    for j in range(1, data[0] + 1):
        graph[i].append(data[j]-1)

1) 선거구를 나누기 : combination 내장 라이브러리를 써서 나누었다. 
1 개이상 n//2 개 이하로 파라미터를 주었는데, 순서가 없는 조합쌍이어서, 절반이상 확인하면 중복이 생긴다. 따라서 1~n 개까지 모두 확인할 필요가 없다. 

index = [i for i in range(n)]
for i in range(1, n//2+1):
    for combi in combinations(index, i):
        # first
        sum1, v1 = bfs(list(combi))
        diff_combi = set(index) - set(combi)
        sum2, v2 = bfs(list(diff_combi))

2) 나눈 선거구가 이어져있는지 확인하기 : bfs를 이용하여 풀었다. 
visited 처리를 하는 방법이 다양하겠지만, 미리 그래프를 초기화 해놓지 않고, 방문한 도시를 집어넣는 방식으로 처리했다. 

def bfs(group):
    queue = deque()
    visited = set()
    queue.append(group[0])
    visited.add(group[0])
    total = 0
    while queue:
        now = queue.popleft()
        total += array[now]
        for next in graph[now]:
            if next in group and next not in visited:
                queue.append(next)
                visited.add(next)

    return total, len(visited)

 

전체코드

import sys
from itertools import combinations
from collections import deque, defaultdict
input = sys.stdin.readline

n = int(input())
INF = int(1e9)
answer = INF # 인구 차이의 최솟값
array = list(map(int, input().split())) # 구역의 인구수
graph = defaultdict(list)

for i in range(n):
    data = list(map(int, input().split()))
    for j in range(1, data[0] + 1):
        graph[i].append(data[j]-1)


def bfs(group):
    queue = deque()
    visited = set()
    queue.append(group[0])
    visited.add(group[0])
    total = 0
    while queue:
        now = queue.popleft()
        total += array[now]
        for next in graph[now]:
            if next in group and next not in visited:
                queue.append(next)
                visited.add(next)

    return total, len(visited)



index = [i for i in range(n)]
for i in range(1, n//2+1):
    for combi in combinations(index, i):
        # first
        sum1, v1 = bfs(list(combi))
        diff_combi = set(index) - set(combi)
        sum2, v2 = bfs(list(diff_combi))

        if v1 + v2 == n:
            answer = min(answer, abs(sum1 - sum2))

if answer == INF:
    print(-1)
else:
    print(answer)