[Python]백준 17141 개리맨더링

dfs 구현 부분이 조금 달라져서 C++ 코드를 참조해서 Python으로 아래처럼 변환했다.

c++
pair <int, int> dfs(int here, int value) {
	visited[here]=1;
    pair <int, int> ret = {1, a[here]}
    for (int there: adj[here]){
    	if (comp[there]!=value) continue:
        if (visited[there]) continue;
        pair <int, int> _temp = dfs(there, value);
        ret.first += _temp.first;
        ret.second += _temp.second;
        }
   	return ret;

Python 으로 구현하면

def dfs(here, value):
	visited[here]=1
    ret = [1, a[here]]
    for there in adj[here]:
    if comp[there]!=value:
    	continue
    if visited[there]:
    	continue
    	_temp = dfs(there, value)
    	ret[0] += _temp[0]
    	ret[1] += _temp[1]
	return ret

어려웠던 문제가 이렇게 dfs 함수만 잘 구현하면 쉬워진다.
다만 이 문제에서 실수하기 쉬운 부분이 두 군데 있었는데 하나는 bit masking을 할 때 모두 0비트로 이루어진 경우와 모두 1비트로 이루어진 경우를 제외하도록 for 문을 약간 수정 해야 했다.

for i in range(1, 1<<n -1): 이렇게 하는 이유는 모두 0비트인 경우와 모두 1비트인 경우는 문제의 조건에서 존재할 수 없는 경우이기 때문이고 두 번째로는 최소한 하나의 0인 구역과 1인 구역이 존재해야지 dfs를 각각에 대해서 한번 씩 돌려볼 수 있기 때문이다. (물론 예외 처리도 할 수 있지만 코테를 할 때는 늘 try문을 쓰지 않는 것이 깔끔한 것 같다)

(참고) 비트마스킹 하는 부분(combination 만들기)

for i in range(1,1<<n-1):
    v = [0] * (n+1)  # 탐색에 사용할 visited 배열
    comp = defaultdict(int)
    size = population = 0 # 각각 다른 주소로 0 assign
    for j in range(n):
        if (i & (1 << j)):
            last_comp1 = j+1
            comp[j+1] = 1
        else:
            last_comp2 = j+1
            comp[j+1] = 0

코드 참고

인프런 C++ 10주완성 코딩테스트 큰돌