[Python 3] BOJ - 2132 "나무 위의 벌레"

 # 문제 링크

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

2132번: 나무 위의 벌레

 

 # 풀이

 개인적으로 이 문제의 핵심은 특정 알고리즘이라기 보다는 문제에 대한 관찰이라고 생각한다.

<그림 1. 예제 입력 1의 경우>

 위 그림 1은 문제에서 제시된 입력을 나타낸 그림이다. 예제 입력 1에서 5개의 구슬을 떨어뜨리면 각각의 순서는 2,4,2,5,2가 된다. 이 때, 떨어지는 순서를 무시하고 각 서브트리(자기 자신도 포함)가 가지고 있는 구슬의 갯수를 세어보자. 그렇게 하면 위 그림 1과 같은 결과가 도출될 것이다.

 예를 들어, 3번 노드의 경우 리프 노드가 아니기에 자기 자신이 구슬의 종착지는 아니지만, 자신의 자식들인 4번과 5번에 각각 구슬이 하나가 있기 때문에 2개가 된다. 4번의 경우, 리프 노드이기에 자기 자신이 구슬의 종착지가 된다. 그러나, 리프 노드이기에 자식이 없으므로 1개가 된다.

<그림 2. 그림 1에서 구슬을 한번 더 떨어뜨린 경우>

 위 그림 2는 위 그림 1에서 구슬 하나를 추가로 떨어뜨린 경우이다. 규칙에 따라, 1번 노드에 도착한 구슬은 왼쪽 자식인 2번 노드 혹은 오른쪽 자식인 3번 노드 중 하나로 떨어져야 한다. 그런데, 그림 1에서 두 노드가 보유한 구슬의 갯수는 같으므로 왼쪽 자식인 2번에 떨어지게 된다. 그런데 2번 노드의 경우 리프 노드이므로 구슬은 멈추게 된다.

<그림 3. 그림 2에서 구슬을 추가로 떨어뜨린 경우>

 위 그림 3은 그림 2에서 다시 한번 구슬 하나를 떨어뜨린 경우이다. 규칙에 따라 구슬은 4번 노드에 도착하게 된다. 만일 여기서 또 하나를 더 떨어뜨린다고 했을 때, 구슬은 어디에 도착해야 할까? 당연히 5번 노드일 것이다. 그 다음은 어떻게 될 것인가? 또 그 다음은?

 한 가지 확실한 것은, 어떤 노드의 왼쪽 자식이 보유한 구슬수는 오른쪽 자식이 보유한 구슬수보다 무조건 1개가 더 많거나 같아야만 한다는 것이다. 위 그림 3을 기준으로 생각해보자. 만일 그림 3에서 떨어뜨린 구슬이 4로 간다고 가정해보자. 현재 그림 3에서 새로운 구슬을 떨어지기 전까지 트리는 문제의 조건을 완벽히 만족한다. 그런데, 새로 떨어뜨린 구슬이 4번 노드로 가게 된다면, 왼쪽 자식이 오른쪽 자식보다 2개 더 많은 구슬을 보유하게 되고, 이는 명백히 조건을 위반한 것이기 때문이다. 추가적으로 만일 이것을 허용하게 되면 루트 노드인 1번 노드의 왼쪽 자식인 2번 노드가 보유한 구슬은 3개, 오른쪽 자식인 3번 노드가 보유한 구슬은 4개가 되어버리므로, 문제의 조건을 맞추기 위해서는 왼쪽 자식인 2번 노드에 최소 구슬이 하나 더 있어야 한다. 그런데, 그렇게 되어버리면 1개의 구슬만을 떨어뜨렸는데 2개의 구슬을 떨어뜨린 결과가 되기 때문에 모순이 발생한다.

 따라서, 전체 떨어뜨린 구슬의 수가 짝수라면, 왼쪽 서브트리와 오른쪽 서브트리가 각각 절반씩을 나누어가지게 되며, 홀수인 경우 왼쪽 서브트리가 오른쪽 서브트리 보다 1개 더 많은 구슬을 보유하게 된다. 이것을 재귀적으로 실행하게 되면 현재 k개의 구슬을 떨어뜨렸을 때 각 노드가 가진 구슬의 수를 알 수 있다.

 그런데 우리가 원한 것은 k번째 떨어뜨린 구슬이 몇번 노드에 있는지이다. 따라서 k-1번째 구슬을 떨어뜨린 결과와 k번째 구슬을 떨어뜨린 결과를 비교하여 각각의 리프 노드가 가진 구슬의 수를 비교하였을 때 이전과 다른 노드가 k번째 떨어뜨린 구슬이 위치하는 곳임은 자명하다.

 

 # 코드

import sys

# 재귀 깊이 해제
sys.setrecursionlimit(300000)

# dfs : 현재 now번째 노드이고 현재 가진 구슬 수가 num개일 때 now번째 노드가 가져야 하는 구슬을
#       temp에 저장하는 함수
def dfs(now, num, temp):
    # 현재 구슬 수 저장
    temp[now] = num
    # Base Case : 리프 노드인 경우
    if tr[now][0] == -1 and tr[now][1] == -1:
        return
    # 오른쪽 자식만 있는 경우
    elif tr[now][0] == -1 and tr[now][1] != -1:
        dfs(tr[now][1], num, temp)
    # 왼쪽 자식만 있는 경우
    elif tr[now][0] != -1 and tr[now][1] == -1:
        dfs(tr[now][0], num, temp)
    # 자식이 두명인 경우
    else:
        dfs(tr[now][0], num - num // 2, temp)
        dfs(tr[now][1], num // 2, temp)

# 입력부
n = int(sys.stdin.readline())
# tr : 현재 idx번째 노드의 왼쪽 자식을 0, 오른쪽 자식을 1로 갖는 2차원 배열
tr = [[-1, -1] for _ in range(n + 1)]
for i in range(1, n + 1):
    a, b = map(int, sys.stdin.readline().split())
    tr[i][0], tr[i][1] = a, b
k = int(sys.stdin.readline())
now = [0] * (n + 1)
before = [0] * (n + 1)
# 각각 k와 k-1번째 구슬을 떨어뜨린 후 비교
dfs(1, k, now)
dfs(1, k - 1, before)
for i in range(1, n + 1):
    if tr[i][0] == -1 and tr[i][1] == -1:
        # 이전과 다른 노드라면 해당 노드가 정답이므로 출력 후 break
        if abs(now[i] - before[i]) == 1:
            print(i)
            break