[HackerRank w7] The crazy helix

문제: HackerRank

이 문제는 초기 배열 [1, 2, 3, ..., n]에서 구간을 flip 하는 연산과 i번째 위치에 있는 수를 구하는 연산, 수 i의 위치를 구하는 연산에 대해 처리 빠르게 처리해야되는 문제다.

Balanced Binary Search Tree라는 것이 있다. 잘 알려진 것으로는 AVL tree, Red-black tree 등등이 있다. 이러한 자료구조로 구현 된 것으로는 STL container의 set과 map이 있다. STL의 응용으로 Balanced BST를 이용해야 풀 수 있는 문제는 거의 존재하지 않는데, 이 문제는 Balanced BST 중 Splay tree라는 자료구조를 이용하면 쉽게 풀린다.

Splay tree에 대한 자세한 설명은 여기를 참고한다. 간단히 소개하자면, splay tree는 splay(i) = i를 BST의 root로 만드는 연산을 통하여 삽입과 삭제를 worst case amortized time complexity $O(\log_2 n)$ 만에 해결한다. 여기서 worst case amortized 라는 것은 (최악인) 데이터에서 여러 연산이 있을 때 이들의 처리 시간 평균을 의미한다. 즉, 한 번의 연산이 $O(n)$이 걸릴 수 있지만, 그 연산을 여러 번 했을 때에는 최악에 연산당 시간복잡도가 $O(\log_2 n)$인 꼴이라는 의미다. 이에 대한 증명은 논문을 참고하자.

이 문제에서 Splay tree를 만드는데 자기보다 위치가 왼쪽에 있는 노드들은 왼쪽 서브트리에, 자기보다 위치가 오른쪽에 있는 노드들은 오른쪽 서브트리에 위치하도록 BST를 만든다. 우리는 splay 연산을 통해, 문제에서 각 연산을 $O(\log_2 n)$ 만에 해결할 것이다.

1) i번째 위치의 값 구하기

이 트리는 BST이다. 왼쪽 서브트리에 있는 노드의 개수를 통해 BST 내에서 탐색과 동일하게 하면 된다.

2) 수 i의 위치 구하기

splay(i) 연산을 할 경우 수 i와 대응하는 노드 i가 루트가 된다. 이 때 왼쪽 서브트리의 크기를 통해 구할 수 있다.

3) 위치 구간 [A,B]를 뒤집기

A-1번째 위치에 있는 노드를 prv, B+1번쨰 위치에 있는 노드를 nxt라 하자. 이는 1) 과 같은 과정을 통해 구할 수 있다.

splay(prv) 연산을 하면, 노드 prv가 루트가 된다. splay(nxt,prv) 연산을 하는데, 이는 nxt를 prv의 오른쪽 자식 노드로 splay하는 연산을 의미한다. 이 연산을 모두 거치면, 루트는 prv이며, 루트의 오른쪽 자식 노드는 nxt이다. 이 때, nxt의 왼쪽 서브트리에 해당하는 노드는 위치 구간 [A,B]에 있는 노드들이며, 그 외 다른 노드들은 없다. 즉, BST에서 어떤 서브트리를 뒤집는 연산을 하면 된다. 서브트리의 루트의 왼쪽 서브트리와 오른쪽 서브트리를 바꾸고, 서브트리의 루트의 왼쪽 서브트리와 오른쪽 서브트리에 대해서 마찬가지의 과정을 적용해주면된다. 매 뒤집는 연산마다 재귀적으로 해주면 $O(n)$이 되므로 그러지말고, 노드별로 뒤집는 여부를 저장해놓고, 1), 2)의 연산이나 splay 연산 처리 중에 뒤집으면 속도가 매우 빠르다.

#include <bits/stdc++.h> using namespace std; #define MAXN 100005 int N, Q, X; struct NODE{ bool flip; int v, sz; NODE *left, *right, *p; void relax(){ if (flip){ swap(left, right); if (left) left->flip ^= 1; if (right) right->flip ^= 1; flip = 0; } } void upd(){ sz = 1; if (left) sz += left->sz; if (right) sz += right->sz; } } *root, *node[MAXN], container[MAXN]; NODE *new_node() { return &container[X++]; } NODE *build_bst(int s, int e) { if (s > e) return NULL; NODE *ret = new_node(); int m = (s+e)>>1; node[m] = ret; ret->v = m; ret->left = build_bst(s, m-1); if (ret->left) ret->left->p = ret; ret->right = build_bst(m+1, e); if (ret->right) ret->right->p = ret; ret->upd(); return ret; } void rotate_right(NODE *p) { NODE *g = p->p, *t = p->left; p->left = t->right; if (t->right) t->right->p = p; t->right = p; p->p = t; if (g){ if (g->left == p) g->left = t; else g->right = t; } t->p = g; p->upd(); t->upd(); } void rotate_left(NODE *p) { NODE *g = p->p, *t = p->right; p->right = t->left; if (t->left) t->left->p = p; t->left = p; p->p = t; if (g){ if (g->left == p) g->left = t; else g->right = t; } t->p = g; p->upd(); t->upd(); } void splay(NODE *t, NODE *rt=0) { while (t->p != rt){ NODE *p = t->p; if (p->p != rt){ NODE *g = p->p; g->relax(); p->relax(); t->relax(); if (g->left == p && p->left == t) rotate_right(g), rotate_right(p); else if (g->left == p && p->right == t) rotate_left(p), rotate_right(g); else if (g->right == p && p->left == t) rotate_right(p), rotate_left(g); else rotate_left(g), rotate_left(p); }else{ p->relax(); t->relax(); if (p->left == t) rotate_right(p); else rotate_left(p); } } if (!rt) root = t; } NODE *find_kth(NODE *now, int k) { now->relax(); if (now->left && k < now->left->sz) return find_kth(now->left, k); if (now->left) k -= now->left->sz; if (!k) return now; return find_kth(now->right, --k); } int main() { scanf("%d%d", &N, &Q); root = build_bst(0, N+1); while (Q--){ int t; scanf("%d", &t); if (t == 1){ // flip int s, e; scanf("%d%d", &s, &e); NODE *p = find_kth(root, s-1); splay(p); NODE *q = find_kth(root, e+1); splay(q, p); q->left->flip ^= 1; }else if (t == 2){ // print pos[x] int x; scanf("%d", &x); splay(node[x]); printf("element %d is at position %d\n", x, root->left->sz); }else if (t == 3){ // print A[x] int x; scanf("%d", &x); NODE *nd = find_kth(root, x); splay(nd); printf("element at position %d is %d\n", x, nd->v); } } }