백준 17471 게리맨더링 (파이썬)

아이디어

• itertools를 사용해 구역의 조합을 구한다.
• 두 구역이 연결되어 있는지 확인해야 한다.

 

---

풀이 및 해설

import sys
from collections import deque, defaultdict
from itertools import combinations
# sys.stdin = open('게리멘더링.txt', 'r')

# 연결이 되어있는지 확인하는 함수
def is_connected(group):
    Q = deque()
    Q.append(group[0])
    visit = [0] * (N+1)
    visit[group[0]] = 1

    while Q:
        u = Q.popleft()
        for v in graph[u]:
            if v in group and visit[v] == 0:
                Q.append(v)
                visit[v] = 1
    for g in group:
        if visit[g] == 0:
            return False
    return True

N = int(sys.stdin.readline().strip())
V_info = [0] + list(map(int, sys.stdin.readline().strip().split()))

graph = defaultdict(list)
# 인접 행렬을 딕셔너리형태로 받는다.
for _ in range(1, N+1):
    tmp = list(map(int, sys.stdin.readline().strip().split()))
    graph[_].extend(tmp[1:])

# 정점들
V = [i for i in range(1, N+1)]
rst = 9e9
# 두 그룹으로 나누기 때문에 N//2+1 이상으로 넘어가면 조합이 중복되어 똑같은 작업을 한다.
for i in range(1, N//2+1):
    # group1에 대해선 itertools사용, group2는 자동적으로 group1을 포함하지 않는 그룹임
    for group1 in combinations(V, i):
        group2 = [j for j in V if j not in group1]
        # 두 그룹에 대해서 연결이 되는 그래프인지 확인을 해봐야된다.
        if is_connected(group1) and is_connected(group2):
            group1_num = sum([V_info[g] for g in group1])
            group2_num = sum([V_info[g] for g in group2])
            rst = min(rst, abs(group1_num-group2_num))
if rst == 9e9:
    print(-1)
else:
    print(rst)

 

---

후기

• 처음 접근의 고민을 많이 했던 문제다.
• 첫 번째 구역이 될 수 있는 조합을 찾은 뒤 그 구역과 나머지 구역이 둘 다 연결되었는지 확인해야했다.
• is_connected함수를 처음에는 dfs로 확인하려고 했는데, 코드가 지저분해지고 조금 복잡해서 bfs로 했다.
• 풀고 보면 상당히 간단한 문제이다.. 근데, 접근 방식을 너무 오래 고민했다.