[C++] 이진 트리 + 힙, 트립

트립 (treap)

 보통 STL에서 제공하는 이진 검색 트리는 대부분 X보다 작은 원소의 수를 계산하거나, k번쩨 원소를 찾는 연산을 지원하지 않는다. 따라서 필요할 때는 직접 군형 잡힌 이진 트리를 구현해야 하는데, 이게 쉽지 않은 작업이다. (속도나 효율도 별로다..) 그렇기에 필요한 경우 간단하게 구현이 가능한 treap을 사용한다.

 

 정의

 트립은 입력이 특정 순서로 주어질 때 성능이 떨어진다는 이진 검색 트리의 단점을 해결하기 위해 고안된 이진 검색 트리이다. 원리는 트리의 형태가 원소들의 추가 순서에 따라 결정되지 않고, 난수에 의해 임의대로 결정된다. 그렇기에 노드가 추가되거나 삭제되더라도 트리 높이의 기대치는 항상 일정하다.

 이와 같은 속성을 만족시키려면 두가지 조건이 성립돼야 한다.

 

1. 이진 검색 트리의 조건 : 모든 노드에 대해 왼쪽 서브트리는 원소 값이 작고, 오른쪽 서브트리는 원소 값이 커야 한다.

2. 힙의 조건 : 모든 노드의 우선순위는 각자의 자식 노드보다 크거나 같다.

 

이름이 treap(= tree + heap)인 이유다.

 

트립

 

 트립 구현

typedef int KeyType;
// 트립의 한 노드를 저장
struct Node
{	
	// 노드에 저장된 원소
	KeyType key;
	int priority, size;

	// 노드 자식
	Node* left, *right;
	
	// 생성자를 통해 난수 우선순위 생성, size와 right/left를 초기화한다.
	Node(const KeyType& _key) :key(_key), priority(rand()), size(1),
		left(NULL), right(NULL) { }
	
	void setLeft(Node* newLeft) { left = newLeft; calcSize(); }
	void setRight(Node* newRight) { left = newRight; calcSize(); }
	void calcSize()
	{
		size = 1;
		if (left) size += left->size;
		if (right) size += right->size;
	}
};

 대게 이진 트리와 다른 점은 노드 객체에서 원소 key 외에도 우선순위 priority(우선 순위)를 가진다는 점이다. 또 다른 점은 size 멤버이다. 이 값은 left, right가 바뀔 때 자동으로 갱신되며, 이 값을 이용하면 k번째 원소를 찾는 연산이나 X보다 작은 원소를 세는 연산 등을 쉽게 구현할 수 있다.

 

 노드의 삽입

 트립에 새 노드 node를 삽입한다고 하자. 가장 먼저 비교할 것은 root와 node의 우선순위이다. 만약 root의 우선순위가 높다면 node는 root아래로 들어가면 된다. 이때 왼쪽이냐 오른쪽이냐는 두 노드의 원소를 비교해서 결정하면 된다. 비교적 과정이 쉽다.

 그러나 root의 우선순위가 node의 우선순위보다 작은 경우 문제가 된다. node가 기존의 root를 밀어내고 새로운 root가 돼야 한다. 이때 사용하는 좋은 방법으로 쪼개기다. 쪼개기는 node를 기준으로 작은 원소만 갖는 서브트리 하나, 큰 원소만 갖는 서브트리 하나를 만드는 과정이다. 이 방법으로 node에 서브트리를 붙여주면 삽입 작업이 완료된다.

 

   노드의 삽입 구현

typedef pair<Node*, Node*> NodePair;
// 기존 트립에서 key 미만 값과 이상 값을 갖는 
// 두 개의 트립으로 분리한다.
NodePair split(Node* root, KeyType key)
{
	if (root == NULL) return NodePair(NULL, NULL);

	if (root->key < key)
	{
		NodePair rs = split(root->right, key);
		root->setRight(rs.first);
		return NodePair(root, rs.second);
	}

	// root가 key 이상인 경우
	NodePair ls = split(root->left, key);
	root->setLeft(ls.second);
	return NodePair(ls.first, root);
}

Node* insert(Node* root, Node* node)
{
	if (root == NULL) return node;

	if (root->priority < node->priority)
	{
		NodePair splitted = split(root, node->key);
		node->setLeft(splitted.first);
		node->setRight(splitted.second);
		return node;
	}
	else if (node->key < root->key)
	{
		root->setLeft(insert(root->left, node));
	}
	else
	{
		root->setRight(insert(root->right, node));
	}
	return root;
}

 

 노드의 삭제

 트립에서의 삭제는 기존 이진 검색 트리에서의 노드 삭제와 별로 다를 것이 없다. 두 서브트리를 합칠 때 어느 쪽이 root가 돼야 하는지를 우선순위로 판단하는 것만 다르다.

   

   노드의 삭제 구현

Node* merge(Node* a, Node* b)
{
	if (a == NULL) return b;
	if (b == NULL) return a;
	if (a->priority < b->priority)
	{
		b->setLeft(merge(a, b->left));
		return b;
	}
	a->setRight(merge(a->right, b));
	return a;
}

Node* erase(Node* root, KeyType key)
{
	if (root == NULL) return root;

	if (root->key == key)
	{
		Node* ret = merge(root->left, root->right);
		delete root;
		return ret;
	}
	if (key < root->key)
	{
		root->setLeft(erase(root->left, key));
	}
	else
	{
		root->setRight(erase(root->right, key));
	}
	return root;
}

 

 k번째 원소 찾기

 굳이 이진 검색 트리를 직접 작성하는 것은 STL에서 제공하지 않은 기능을 사용할 때이다. 그중 하나가 트리를 크기 순으로 나열했을 때,  k번째에 오는 노드를 찾을 경우이다. 트립인 경우 size를 계산해서 저장해두기에 이것을 구현하기가 어렵지 않다.

 root, root->left, root->right가 있을 때, 각 서브트리의 크기를 알고 있으면 k번째 노드가 어디에 있는지 알 수 있다. 왼쪽 서브트리의 크기가 l일 때, 다음과 같이 세 개의 경우로 나눌 수 있다.

 

1. k <= l : k번째 노드는 왼쪽 서브트리에 있다.

2. k = l+1 : 루트가 k번째 노드이다.

3. k > l+1 : k번째 노드는 오른쪽 서브트리에서 k-l-1번째 노드가 된다.

 

   k번째 원소 찾기 구현

Node* kth(Node* root, int k)
{
	int leftSize = 0;
	if (root->left != NULL) leftSize = root->left->size;
	if (k <= leftSize) return kth(root->left, k);
	if (k == leftSize + 1) return root;
	return kth(root->right, k - leftSize - 1);
}

 

 X보다 작은 원소 세기

 또 다른 유용한 연산으로 특정 범위 [a, b)가 주어질 때 범위 안에 들어가는 원소들의 숫자를 계산하는 연산이 있다. 만약, root의 원소가 X와 같거나 더 크면 root와 그 오른쪽에 있는 서브트리에 있는 원소들은 모두 X 이상이므로, 왼쪽 서브트리에 있는 원소들만 재귀적으로 세서 반환하면 된다. 만약 root의 원소가 X보다 작다면, 그 왼쪽 서브트리에 있는 원소들은 모드 X보다 작으므로 재귀적으로 세지 않아도 전부 답에 들언다는 것을 알 수 있다. 오른쪽 서브트레이서 X보다 작은 원소들의 재귀적으로 찾고, 이것에 왼쪽 서브트리에 포함된 노드 및 루트의 개수를 더해주면 된다.

 

   X보다 작은 원소 세기 구현

int countLessThan(Node* root, KeyType key)
{
	if (root == NULL) return 0;
	if (root->key >= key)
	{
		return countLessThan(root->left, key);
	}
	int ls = (root->left ? root->left->size : 0);
	return ls + 1 + countLessThan(root->right, key);
}