Treapの問題をC++で解く

アルゴリズムとデータ構造入門

当記事について
ALDS1_8_D: Treap: Treap
1. 考えたこと
2. コード例
おわり

当記事について

Aizu Online Judgeに掲載されている問題をC++で解くという遊びをしているので、考え方の例と解けたコード例を記録していこうかと思います。コード例はたたんでおこうかと思います。クリックやタップでコード例を開けます。

問題はこちら:

Aizu Online Judge

ALDS1_8_D: Treap: Treap

二分探索木とヒープを融合させたものを実装する。乱択平衡二分木の一種。

乱数とキーは与えられるので、指示通りに実装していけば問題なくできると思います。

手前が二分探索木関連の問題のようなので、やっておくと簡単かもしれません。

考えたこと

標準入力で文字列を受け取るので、<string>ヘッダをインクルードしておきます。

オブジェクト指向でクラスを利用して実装しようかとも思いましたが、あまり必要性がないかなと思い例通りの関数で実装しました。

Treapというものを知らなかったので回転は実際に過程を図を書いて理解しました。

ノードは構造体で定義して、木なのでルートだけ保持します。

ノードの探索は、keyを見て、自身を返すか左右のどちらかを再帰的に呼び出せばいいです。

中間順巡回と先行順巡回は確か以前の問題でやったと思います。このワードで調べれば該当問題か解説が真っ先にヒットすると思います。

簡単な説明: 例えば、あるノードAの左の子をL、右の子をRとすると、中間順巡回は、L --> A --> Rの順で巡回し、先行順巡回は、A --> L --> Rの順で巡回します。

コード例

クリックまたはタップで開閉できます

// Treap

#include <iostream>
#include <string>

struct Node
{
    int key, priority;
    Node *left, *right;
    Node(int key, int priority)
    {
        this->key = key;
        this->priority = priority;
        left = right = nullptr;
    }
};

Node *rightRotate(Node *t);
Node *leftRotate(Node *t);
Node *insert(Node *t, int key, int priority);
Node *remove(Node *t, int key);
Node *remove_target(Node *t, int key);
Node *find(Node *t, int key);
// print inorder tree walk
void print_itw(Node *t);
// print preorder tree walk
void print_ptw(Node *t);

int main()
{
    Node *root = nullptr;
    int m;
    std::cin >> m;
    for (size_t i = 0; i < m; i++)
    {
        int k, p;
        std::string op;
        std::cin >> op;
        if (op == "insert")
        {
            std::cin >> k >> p;
            root = insert(root, k, p);
        }
        else if (op == "find")
        {
            std::cin >> k;
            if (find(root, k) != nullptr)
                std::cout << "yes" << std::endl;
            else
                std::cout << "no" << std::endl;
        }
        else if (op == "delete")
        {
            std::cin >> k;
            root = remove(root, k);
        }
        else if (op == "print")
        {
            print_itw(root);
            std::cout << std::endl;
            print_ptw(root);
            std::cout << std::endl;
        }
    }
}

Node *rightRotate(Node *t)
{
    Node *s = t->left;
    t->left = s->right;
    s->right = t;
    return s;
}

Node *leftRotate(Node *t)
{
    Node *s = t->right;
    t->right = s->left;
    s->left = t;
    return s;
}

Node *insert(Node *t, int key, int priority)
{
    if (t == nullptr)
        return new Node(key, priority);
    if (key == t->key)
        return t;

    if (key < t->key)
    {
        t->left = insert(t->left, key, priority);
        if (t->priority < t->left->priority)
            t = rightRotate(t);
    }
    else
    {
        t->right = insert(t->right, key, priority);
        if (t->priority < t->right->priority)
            t = leftRotate(t);
    }
    return t;
}

Node *remove(Node *t, int key)
{
    if (t == nullptr)
        return nullptr;

    if (key < t->key)
        t->left = remove(t->left, key);
    else if (key > t->key)
        t->right = remove(t->right, key);
    else
        return remove_target(t, key);
    return t;
}

Node *remove_target(Node *t, int key)
{
    if (t->left == nullptr && t->right == nullptr)
    {
        delete t;
        return nullptr;
    }
    else if (t->left == nullptr)
        t = leftRotate(t);
    else if (t->right == nullptr)
        t = rightRotate(t);
    else
    {
        if (t->left->priority > t->right->priority)
            t = rightRotate(t);
        else
            t = leftRotate(t);
    }
    return remove(t, key);
}

Node *find(Node *t, int key)
{
    if (t == nullptr)
        return nullptr;
    if (t->key == key)
        return t;
    else if (key < t->key)
        return find(t->left, key);
    else
        return find(t->right, key);
}

void print_itw(Node *t)
{
    if (t == nullptr)
        return;
    print_itw(t->left);
    std::cout << " " << t->key;
    print_itw(t->right);
}

void print_ptw(Node *t)
{
    if (t == nullptr)
        return;
    std::cout << " " << t->key;
    print_ptw(t->left);
    print_ptw(t->right);
}