[백준 2957번] 이진 탐색 트리 (Python3)

import sys
input = sys.stdin.readline

def update(i):
  check[i] = 1
  while i<=M:
    fen[i] += 1
    i += i&-i

def cal(i):
  S = 0
  while i:
    S += fen[i]
    i -= i&-i
  return S

def find(x,i,c):
  s = cal(i)
  if s==x and check[i]:
    return i
  if s<x:
    return find(x,i+(1<<c),c-1)
  return find(x,i-(1<<c),c-1)

N = int(input()); M = 1<<19
fen,check,depth = [[0]*(M+1) for i in range(3)]
C = 0
for _ in range(N):
  x = int(input()); s = cal(x)
  if fen[M]==0:
    pass
  elif s==0:
    depth[x] = depth[find(1,M,18)]+1
  elif s==fen[M]:
    depth[x] = depth[find(s,M,18)]+1
  else:
    depth[x] = max(depth[find(s,M,18)],depth[find(s+1,M,18)])+1
  C += depth[x]
  update(x)
  print(C)

팬윅트리로 구현하였다.

 

당연하게도 이 문제를 지문에 나온 코드를 그대로 구현하면 시간초과가 발생하게 된다. 단순하게 생각해서 입력값이 1~N까지 오름차순으로 주어질 경우 시간복잡도가 O(N^2)이 되기 때문이다. 따라서 문제를 해결하기 위해서는 아이디어가 필요하다.

이진탐색트리를 관찰해보면, 어떤 값 x를 삽입할 때, x의 깊이는 현재까지 삽입한 모든 값 중에서 x보다 작은 값 중 가장 큰 값을 left, x보다 큰 값 중 가장 작은 값을 right라고 한다면, left와 right 사이의 공간에 삽입된다는 것을 알 수 있다. 또 부모노드는 left와 right 중 하나가 되며, 이때 깊이가 더 깊은 값이 x의 부모노드가 된다.

즉, 문제의 핵심은 left와 right 값을 찾는것이 될 것이다.

 

풀이를 설명하면,

1. 값이 입력될 때마다 left와 right를 찾는다. 찾는 방식은 팬윅트리(혹은 세그먼트 트리)를 이용한다. (방법은 [백준 2243번] 사탕상자 (Python3) (tistory.com)와 비슷)

2. 현재 값의 깊이는 max(depth[left],depth[right])+1이 된다.

3. C에 현재 값의 깊이를 더하고 출력한다.

 

문제를 풀고 나서 다른 사람들 풀이를 확인했는데, C++로 해결한 경우가 많았고, C++의 경우 lower_bound, upper_bound라는 함수를 통해서 1번 과정을 매우 쉽게 해결했다는 것을 알게되었다... 저 함수를 이용하면 세그먼트 트리나 팬윅트리 없이고 left와 right값을 한번에 구할 수 있다.........

문제 제출자가 파이썬은 고려하지 않고 낸 것 같다.

 

 

오늘의 교훈) 파이썬 화이팅