균형잡힌 이진검색 트리 (TREAP)

PS에서 사용하기 좋은 균형잡힌 이진검색트리 treap. 구현이 상대적으로 간편하다

각 노드마다 임의의 우선순위를 가지도록 한다. 그리고 이진검색트리 성질을 유지하면서 우선순위가 높은 노드가 무조건 위에 오도록 잘 구현하면, 아주 높은 확률(?)로 balanced됨을 증명할 수 있다. 증명이 중요한게 아니라 구현이 중요하다. 어떻게 이진검색트리의 성질을 유지하면서 우선순위규칙을 지키도록 구현할 수 있을까?

구현

Node 구조체

typedef int KeyType;
struct Node {
    KeyType key;
    int priority, size;
    Node* left, * right;

    Node(const KeyType& _key) : key(_key), priority(rand()),
        size(1), left(NULL), right(NULL) {}
    // 재귀적으로 루트부터 전부 free할 수 있음
    ~Node() {
        delete left;
        delete right;
    }

    void setLeft(Node* newLeft) {
        left = newLeft;
        calcSize();
    }
    void setRight(Node* newRight) {
        right = newRight;
        calcSize();
    }
    void calcSize() {
        size = 1;
        if (left) size += left->size;
        if (right) size += right->size;
    }
};

삽입

typedef pair<Node*, Node*> NodePair;
NodePair split(Node* root, KeyType key) {
    if (root == NULL) return NodePair(NULL, NULL);

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

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

    // 우선순위 앞설경우, root를 둘로나누고 
    // node 좌우에 붙이고 node가 새로운 루트
    if (root->priority < node->priority) {
        NodePair splitted = split(root, node->key);
        node->setLeft(splitted.first);
        node->setRight(splitted.second);
        return node;
    }
    else if (root->key < node->key)
        root->setRight(insert(root->right, node));
    else
        root->setLeft(insert(root->left, node));

    return root;
}

삭제

// left와 right를 하나로 합쳐서 루트 리턴
Node* merge(Node* left, Node* right) {
    if (left == NULL) return right;
    if (right == NULL) return left;

    if (left->priority < right->priority) {
        right->setLeft(merge(left, right->left));
        return right;
    }
    else {
        left->setRight(merge(left->right, right));
        return left;
    }
}
// 삭제하고 새로운 루트 리턴
Node* erase(Node* root, KeyType key) {
    // 삭제하려는 원소가 없음
    if (root == NULL)
        return root;

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

    if (key < root->key)
        root->setLeft(erase(root->left, key));
    else
        root->setRight(erase(root->right, key));
    return root;
}

사용방법

// 초기화
Node* treap = NULL;
// 삽입
treap = insert(treap, new Node(key));
// 삭제
treap = erase(treap, key);
// 사용후에는 꼭 메모리 할당 해제
// 루트에서만 할당해제하면 재귀적으로 전부 free되도록
// destructor를 적당히 오버라이딩 해놓았다.
delete treap;

활용

k번째 노드 찾기

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

k보다 작은 원소 개수 찾기

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

Reference

• 종만북