[백준/14570] 나무 위의 구슬 (Python)

Problem : https://www.acmicpc.net/problem/14570

14570번: 나무 위의 구슬

 

Difficulty : Gold 1

 

Status : Solved

 

Time : 00:36:36

 

---

 

문제 설명

 

 

 

이진 트리란, 위처럼 모든 노드의 자식의 수가 2개 이하인 트리이다.

각 노드에 쓰여 있는 수는 노드의 번호를 의미한다.

특히, 이 문제에서는 루트가 고정되어 있으며, 노드의 순서가 중요한(어떤 서브트리에서도 좌우를 변경할 수 없는) 이진 트리에 대해 다루기로 한다.

이진 트리의 루트에 구슬을 하나 올려놓으면 구슬은 아래와 같은 과정을 거쳐 떨어진다.

 

1. 현재 구슬이 놓인 노드의 자식이 없다면 그 자리에서 멈춘다.
2. 1을 만족하지 않으며, 만일 현재 구슬이 놓인 노드의 자식 노드가 한 개라면 해당 자식 노드로 떨어진다.
3. 1, 2를 만족하지 않으며, 만일 현재 구슬이 놓인 노드의 자식 노드가 두 개라면, 현재 노드의 왼쪽 서브트리에 담긴 모든 구슬의 수 <= 오른쪽 서브트리에 담긴 모든 구슬의 수일 경우, 왼쪽 자식 노드로 떨어진다.
4. 그 외의 경우에는 오른쪽 자식 노드로 떨어진다.
5. 1~4번의 조건을 다시 체크하고 되풀이한다.

구슬은 위와 같은 과정을 거쳐 결국 단말 노드에 쌓이게 된다.

예를 들어, 위의 그림과 같은 트리에 구슬을 떨어뜨릴 경우.
첫 다섯 개의 구슬은 2번, 4번, 2번, 5번, 2번 노드에 차례대로 떨어지게 된다.
위처럼 트리가 충분히 작거나 구슬의 수가 충분히 적을 경우엔 직접 시뮬레이션을 통해
구슬이 떨어지는 순서를 유추할 수가 있다.
하지만, 우리가 관심있는 것은 큰 트리에서 많은 수의 구슬을 떨어뜨리는 과정이다.
임의의 이진 트리가 주어지고, K가 주어졌을 때
K번째 구슬이 어느 노드에서 멈추게 될 지 충분히 빠르게 계산해낼 수 있을까?

 

 

입력 및 출력

 

입력

 

첫 줄에 이진 트리의 노드의 수 N이 주어진다. (1 ≤ N ≤ 200000)
둘째 줄부터 N개의 줄에 걸쳐, U V가 주어진다.
i번째 줄에 주어지는 U, V는 각각 i번 노드의 왼쪽 자식이 U, 오른쪽 자식이 V임을 의미한다.
만약 U = -1 또는 V = -1이라면, 해당 위치에 자식 노드가 존재하지 않는다는 것이다.
그 외의 경우엔 항상 2 ≤ U, V ≤ N을 만족한다.
이어 마지막 줄에 문제에서 설명한 K가 주어진다. (1 ≤ K ≤ 1018)
주어지는 트리는 항상 올바른 이진 트리임이 보장되며, 루트는 항상 1번 노드이다.

 

출력

 

K번째 구슬이 떨어지는 노드의 번호를 출력한다.

 

입력 예시

 

5
2 3
-1 -1
4 5
-1 -1
-1 -1
5

 

출력 예시

 

2

 

 

---

 

풀이

 

일일히 K번 브루트포스를 시행하면 시간 초과(O(N*K)), 패턴을 직접 찾아 시뮬레이션하는 경우는 공간복잡도 초과가 일어난다. 즉 접근을 달리해볼 필요가 있다.

 

임의의 서브트리에 k번째 수를 체크하고 싶다.

• 만약 리프 노드라면, k번 반복하더라도 자기 자신이므로 자기 자신을 반환한다.
• 만약 자식이 하나뿐이라면, 자식 노드 단에서 판별이 가능하다.
• 만약 자식이 둘이라면, 왼쪽, 오른쪽, 왼쪽, 오른쪽... 순으로 두 서브트리에 교차하여 구슬을 흘리게 된다.
  • 즉 왼쪽 서브트리에 k // 2번 패턴이 진행되고, 오른쪽 서브트리에도 k // 2번 패턴이 진행될 것이다.
  • 여기서 k % 2 == 0이라면 k번째 패턴은 왼쪽 서브트리상에 존재한다. 또한 왼쪽 서브트리에서 k // 2번 진행했으므로, k // 2번으로 갱신한 후 반환할 수 있다. 반면 k % 2 == 1이라면 마찬가지로 k번째 패턴이 오른쪽 서브트리상에 k // 2번째 패턴으로 존재한다. 이를 반환하도록 한다.

 

여러모로 트리 패턴 찾기가 겹쳐서 초반에 해멨던 것 같다. 이 문제도 함께 풀어보자.

 

2023.02.05 - [PS/프로그래머스] - [프로그래머스] 1, 2, 3 떨어뜨리기 (Python)

[프로그래머스] 1, 2, 3 떨어뜨리기 (Python)

 

풀이 코드

import sys
input = sys.stdin.readline
sys.setrecursionlimit(300000)

N = int(input())
children = [[0, 0]] + [list(map(int, input().split())) for _ in range(N)]
K = int(input())

def dfs(node, k) :
  if children[node][0] == children[node][1] == -1 :
    return node
  if children[node][0] == -1 :
    return dfs(children[node][1], k)
  if children[node][1] == -1 :
    return dfs(children[node][0], k)

  if k % 2 == 0 :
    return dfs(children[node][0], k // 2)
  return dfs(children[node][1], k // 2)

print(dfs(1, K-1))

 

풀이하다 보니 이 문제는 재귀 형태가 아닌 반복문 형태가 더 적합할 듯 싶어 코드를 수정했다. 이 경우도 정상적으로 동작한다!

 

import sys
input = sys.stdin.readline

N = int(input())
children = [[0, 0]] + [list(map(int, input().split())) for _ in range(N)]
K = int(input()) - 1

node = 1
while True :
  if children[node][0] == children[node][1] == -1 :
    print(node)
    break
  if children[node][0] == -1 :
    node = children[node][1]
  elif children[node][1] == -1 :
    node = children[node][0]
  else :
    node, K = children[node][K % 2], K // 2

풀이 완료!