[BOJ] 게리맨더링 in Python

박준규·2022년 5월 18일

알고리즘

목록 보기

33/39

의식의 흐름

예전에 풀어본 문제지만, 그때는 안풀려서 다른 분 답을 참고했다. bitmask를 활용한 알고리즘 문제를 꼭 내 힘으로 풀어보겠다는 다짐을 한체 꾸준히 bitmask를 풀었고 어느 정도 감이 잡혔던 상태로 다시 한 번 이 문제에 도전했다. 문제는 그렇게 어렵지 않다. 하지만 이것을 어떻게 bitmask로 옮길 것인가? 는 좀 다른 질문인것 같다.

여튼 문제를 꼼꼼히 읽고 어떻게 해야 할지 설계를 했다.

문제의 노드의 개수는 10개로 입력값이 많지 않다. 따라서 bitmask를 활용한 완전 탐색을 진행하는 것이 효율적이다.

일단 각각의 노드가 존재한다는 점
각각의 노드로 2개의 집단을 나눴을 때 꼭 집단 안에서 모든 노드를 방문할 수 있어야 한다는 점
모든 노드는 집단에 있어야 한다는 점
각 집단의 state를 bitmask로 나눌 수 있다는 점
문제에서 주어진 데이터는 graph 형식으로 다룰 수 있다는 점

그래서 일단 데이터를 받고 각각의 상태를 구분했다.

이를 위해 모든 노드를 방문하는 상태를 만들고 특정 상태를 만들고 이에 반대되는 상태를 만들었다. 그게 아래의 코드다.

# 모든 노드를 방문하는 상태
total_state = 1
for i in range(n):
    total_state = total_state | (1 << i)

# group1의 상태와 group2의 상태
for i in range(1 << n):
    state1 = i
    state2 = i^total_state

위 코드를 통해 각각의 상태를 구분했으며, 이 상태들을 통해 방문 여부를 체크해야 한다.

나는 DFS를 이용해서 방문했고, BFS로 방문해도 상관없다. 본인이 편한 탐색 알고리즘을 사용하면 된다.

DFS

def dfs(start, state):
    for nxt in graph[start]:
        if not visited[nxt] and (1 << nxt-1) & state:
            visited[nxt] = True
            dfs(nxt, state)

참고로 다음 노드를 의미하는 nxt에 -1을 한 이유는 1번 노드는 1 << 0 이고 2번 노드는 1 << 1을 의미하기 때문이다.

여튼 위 dfs를 이용하여 다음과 같이 group1, 2의 상태에 방문 여부를 체크한다.

방문 로직

    for j in range(n):
        if state1 & (1 << j) and not visited[j]:
            visited[j+1] = True
            dfs(j+1, state1)
            break
    
    for j in range(n):
        if state2 & (1 << j) and not visited[j]:
            visited[j+1] = True
            dfs(j+1, state2)
            break

솔직히 같은 코드를 작성해서 별로 마음에 들진 않지만, 2개의 집단 밖에 없기 때문에 충분히 문제에 알맞은 코드로 판단했다.

group의 상태를 와 각 비트를 체크하면서 1 << j가 state에 속하면서 방문하지 않은 노드라면 dfs 순회를 시도하되, 한 번 순회하면 끝내라는 의미이다. 문제의 조건에서 모든 그룹의 노드는 단 한 번의 순회로 모두 방문할 수 있기 때문에 이렇게 진행했다.

그리고 이렇게 순회했을 때 visited는 모두 방문체크가 되어 있어야 한다.

방문 체크

    for j in range(1, n+1):
        if not visited[j]:
            break

만약 위 code에서 break 문제 걸리지 않았다면, 문제에서 주어진 조건을 만족하는 state이기 때문에 다음과 같은 연산을 시도한다.

답 도출

	p1 = 0
    p2 = 0

	for k in range(n):
    	if (1 << k) & state1:
        	p1 += populations[k]
        else:
        	p2 += populations[k]

        answer = min(answer, abs(p1-p2))

전체 코드

import sys
input = sys.stdin.readline

n = int(input())
populations = list(map(int, input().split()))

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

def dfs(start, state):
    for nxt in graph[start]:
        if not visited[nxt] and (1 << nxt-1) & state:
            visited[nxt] = True
            dfs(nxt, state)

total_state = 1
for i in range(n):
    total_state = total_state | (1 << i)

answer = 1001
for i in range(1 << n):

    state1 = i
    state2 = i^total_state
    
    visited = [False for _ in range(n+1)]

    for j in range(n):
        if state1 & (1 << j) and not visited[j]:
            visited[j+1] = True
            dfs(j+1, state1)
            break
    
    for j in range(n):
        if state2 & (1 << j) and not visited[j]:
            visited[j+1] = True
            dfs(j+1, state2)
            break

    for j in range(1, n+1):
        if not visited[j]:
            break
    else:
        p1 = 0
        p2 = 0
        
        for k in range(n):
            if (1 << k) & state1:
                p1 += populations[k]
            else:
                p2 += populations[k]

        answer = min(answer, abs(p1-p2))

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