BFS - 17471번: 게리맨더링

from itertools import combinations
from collections import deque

n = int(input())
line = "0 " + input()
# 인구 수 저장 리스트
amount_list = list(map(int, line.split()))
area_list = []
min_diff = -1
graph = {}

for i in range(1, n + 1):
    area_list.append(i)
    line = list(map(int, input().split()))
    for j in range(1, len(line)):
        if i in graph.keys():
            graph[i].append(line[j])
        else:
            graph[i] = [line[j]]


# 그래프 탐색
def bfs(graph, node, visited, union):
    visited[node] = True
    union.append(node)
    queue = deque([node])

    while queue:
        x = queue.popleft()
        childs = graph.get(x, [])

        for c in childs:
            if not visited[c]:
                visited[c] = True
                union.append(c)
                queue.append(c)
    return


# 그래프 탐색2
def bfs_2(graph, node, visited, union, area):
    visited[node] = True
    union.append(node)
    queue = deque([node])

    while queue:
        x = queue.popleft()
        childs = graph.get(x, [])

        for c in childs:
            # 아직 방문하지 않았고 해당 영역의 노드라면 방문하기
            if not visited[c] and c in area:
                visited[c] = True
                union.append(c)
                queue.append(c)
    return


visited = [False] * (n + 1)
union_list = []  # 존재하는 그래프 리스트

for i in range(1, len(area_list) + 1):
    if not visited[i]:
        union = []
        bfs(graph, i, visited, union)
        union_list.append(union)


# 그래프가 3개 이상이면 2개로 나누는 것은 불가능
if len(union_list) >= 3:
    print("-1")
# 그래프가 2개면 해당 그래프의 인구 수 차이 리턴
elif len(union_list) == 2:
    sum_1 = 0
    for elem in union_list[0]:
        sum_1 += amount_list[elem]
    sum_2 = 0
    for elem in union_list[1]:
        sum_2 += amount_list[elem]
    print(abs(sum_1 - sum_2))
# 그래프가 1개면 진행
elif len(union_list) == 1:
    for k in range(1, n):  # (1 ~ n-1)
        # k개의 가능한 모든 조합을 생성, 해당 조합은 A영역으로 지정하고 나머지는 B영역으로 지정해서 실험
        combies = list(combinations(area_list, k))

        for j in range(0, len(combies)):
            A = list(combies[j])
            B = []
            # A를 제외한 모든 노드를 B에 넣기기
            for i in range(1, n + 1):  # (1 ~ n)
                if i not in A:
                    B.append(i)
            # A와 B가 각각 연결된 그래프인지 확인
            visited_A = [False] * (n + 1)
            visited_B = [False] * (n + 1)
            union_A, union_B = [], []
            bfs_2(graph, A[0], visited_A, union_A, A)
            bfs_2(graph, B[0], visited_B, union_B, B)
            A_realset = set(A)
            B_realset = set(B)
            A_set = set(union_A)
            B_set = set(union_B)

            # 탐색한 union_A와 A가 동일하고, 탐색한 union_B와 B가 동일하다면 각각의 그래프가 서로 연결되어있는 것
            if A_realset == A_set and B_realset == B_set:
                sum_1 = 0
                for elem in union_A:
                    sum_1 += amount_list[elem]
                sum_2 = 0
                for elem in union_B:
                    sum_2 += amount_list[elem]
                # 두 그래프의 인구수 차이 계산
                diff = abs(sum_1 - sum_2)
                # 최소값 갱신
                if min_diff == -1 or min_diff > diff:
                    min_diff = diff
    print(min_diff)


# 그래프를 두개로 쪼개면서 각 노드에 속한 인구수 차이가 최소가 되도록 만들고, 그 때의 최솟값을 리턴턴
# 모든 경우 탐색