Implementing Persistent BST with Weight-Balanced Tree

Splay tree 등의 BBST 구현을 여러번 시도해 보았지만 내가 대회에서 직접 짜면 버그가 너무 많아서 보통 피해왔었다. 얼마 전 QOJ를 읽다가 Weight-Balanced Tree라는 Persistent BBST의 굉장히 간단한 구현체를 찾아 이 블로그에 소개한다. 이 구현의 장점은:

• Rotate 연산을 사용하지 않는다.
  • Rotate 연산은 유도하기 대단히 귀찮아서 보통 암기해야 하고 실수하기 쉽다.
  • 사실 이 구현도 Merge에서 암기해야 할 부분이 있는데 Rotate보단 쉬운 거 같다. 그리고 그 부분은 실수를 조금 해도 괜찮다.
• 리프 노드만이 실제 원소에 대응된다. 리프 노드가 아니면 대응되는 원소는 없으며, 왼쪽 자식 / 오른쪽 자식을 모두 가진다. 일반적으로 우리가 사용하는 세그먼트 트리등과 동일하다.
  • 일반적인 BBST에서는 모든 노드가 실제 데이터에 대응된다는 식으로 구현한다. 고로 각 노드가 왼쪽/오른쪽 자식을 가질 수 있을 수도 있고 없을 수도 있다. 이러한 구현 방식은 메모리를 2배 절약하지만, NULL pointer dereference의 여지가 매우 많다.
  • BBST가 제대로 구현하기 까다롭다는 인상이 이러한 "구현 습관" 때문이 아닐까 생각이 들었다. 예를 들어, Splay tree는 splay 연산 때문에 무조건 각 원소와 노드 사이에 일대일 대응이 있어야 한다.
• Persistent하게 구현할 수 있다.
  • 예를 들어 Splay tree는 amortization을 사용하기에 Persistent한 구현이 불가능하다.
  • Treap도 Persistent한 구현을 하면 $O(\log n)$ 연산이 안 된다고 들었다. 그게 $O(\log^2n)$ 인지 $O(n)$ 인지는 기억이 안 나고, 사실 그 이유도 자명하지 않아서 기억이 안 난다. 특정 연산에서는 될 수도 있다. (사실 그냥 될 수도 있다.)
  • 이 구현은 Persistent하게 구현하기가 아주 쉽고, 다른 Persistent BBST에 비해서 시간 / 메모리 면에서 이점이 있다.

이 구현은 Weight-Balanced Tree라는 BBST를 사용한다. Weight-Balanced Tree는, 모든 노드에 대해서 왼쪽 서브트리의 크기와 오른쪽 서브트리의 크기 비율이 $\alpha$ 이라는 상수 이하라는 조건을 만족하는 트리이다. $\alpha$ 가 상수라면, Weight-Balanced Tree의 모든 노드는 깊이가 $O(\log n)$ 이고 고로 BBST이다.

BBST를 구현하기 위해서는 다음 두 연산만 지원하면 된다:

• Split: 트리 $T$가 있고 정수 $0 \le k \le |T|$ 가 주어질 때, $T$ 의 크기 $k$ prefix와 크기 $|T| - k$ suffix에 대응되는 두 트리를 반환한다.
• Merge: 두 트리 $T_1, T_2$ 가 주어질 때, $T_1$ 과 $T_2$ 의 concatenation을 반환한다.

왜냐면, 나머지 연산이 자연스럽게 따라오기 때문이다:

• Insert: 삽입 위치를 찾고 그 위치 기준으로 Split, 이후 왼쪽 트리 + 삽입 원소 + 오른쪽 트리 Merge.
• Delete: 해당 위치 앞/뒤로 split후 merge.
• 구간 쿼리도 Split으로 구현할 수 있다. (전체 트리가 필요한게 아니면 그냥 세그트리처럼 짜도 됨.)

두 Weight-Balanced Tree를 Merge해 보자. 만약 두 트리의 크기 차이가 $\alpha$ 이하면, 그냥 루트 노드를 하나 만들어주면 되니 자명하다. 일반성을 잃지 않고 $|L| > \alpha |R|$ 라 하자. 첫 번째로, $L = (L_1, L_2)$ 라고 했을 때, 재귀적으로 $(L_2, R)$ 을 merge한다. $X = L_1, Y = merge(L_2, R)$ 이라고 하자.

이제 Weight-Balanced 조건을 만족하게 하면서 $X$ 와 $Y$ 를 Merge해야 한다. 결론적으로, 다음 세 가지 중 하나가 Weight-Balanced Tree를 준다. $X = a, Y = ((b,c),d)$ 라고 하면:

• $(a, ((b, c), d))$ (그냥 (X, Y) 다.)
• $((a, (b, c)), d)$
• $((a, b), (c, d))$

생각해보면, $X$ 와 $Y$ 는 이미 almost balanced되어 있다. $L$ 이 원래 weight-balanced 되어 있으니까, $L$ 의 반쪽을 $R$ 에 줬다고 하면, $R$ 의 원래 비율 자체가 $L$ 에 비해서 작으니 여전히 weight-balanced일 것이다. 그래서 정말 "대충 하면" 되는데, 이걸 증명하려면 생각보다 많이 귀찮다. 다행이도, 위 구현을 작성해 주신 yosupo 님이 상용 계산 프로그램 (Wolfram Engine) 으로 돌려본 결과 $\alpha \geq 1 + \sqrt 2$ 면 항상 저 셋 중 하나가 올바른 weight-balanced tree임을 확인할 수 있었다고 한다. 그래서 여러분은 걱정하지 않고, 그냥 저 세 패턴만 외워서 간단한 Persistent BST를 즐기면 된다.

위 방법으로 내가 직접 구현한 BOJ 13159. 배열 문제의 코드이다: 링크 (거의) Persistent하게 구현했지만 속도 면에서 큰 차이가 없음을 볼 수 있다.

위 방법으로 내가 직접 구현한 BOJ 17486. 수열과 쿼리 30 문제의 코드이다: 링크. 다른 코드 보다 속도 면에서 아주 빠르고, 특히 Garbage collection을 구현하지 않아도 구현을 세심하게 하는 정도에서 메모리 제한 안에 넣을 수 있었다.

어차피 그렇게 robust하게 짜지는 않았으니 첨부한 코드는 참고용으로만 사용하고 실제 문제 해결은 직접 구현해 보았으면 좋겠다. (구체적으로, BOJ 17486 문제에 대해서 위 코드와 지나치게 유사한 코드가 제출될 경우 BOJ 계정 징계가 있을 수 있다.)