指针 AVL 40pts 求调...
查看原帖
指针 AVL 40pts 求调...
655192
Tibrella楼主2023/5/4 19:21

调了一下午了,目前来看插入不能完全保证平衡,但是我想不到怎么做得更好了...

#include <iostream>

auto max = [](auto a, auto b) { return a > b ? a : b; };
auto min = [](auto a, auto b) { return a < b ? a : b; };

template <class Type, int Size>
class AVLTree {
    struct Node {
        Node *lc, *rc;
        int height, siz;
        Type val;
    } tree[Size], *root, *null, *tot, *rubbish_bin[Size];
    int bintop;

    using nodePointer = Node*;

    int get_height(nodePointer nod) {
        if (nod == null) return 0;
        return nod->height;
    }
    int get_siz(nodePointer nod) {
        if (nod == null) return 0;
        return get_siz(nod->lc) + get_siz(nod->rc) + 1;
    }

    void update(nodePointer nod) {  // 此处高度只和下面有关,所以直接获取儿子高度进行更新即可。
        nod->height = max(get_height(nod->lc), get_height(nod->rc)) + 1;
        nod->siz = get_siz(nod->lc) + get_siz(nod->rc) + 1;
    }

    nodePointer left_rotate(nodePointer nod) {
        nodePointer newRoot = nod->rc;
        nod->rc = newRoot->lc;
        newRoot->lc = nod;

        update(nod);
        update(newRoot);

        return newRoot;  // 返回新根
    }

    nodePointer right_rotate(nodePointer nod) {
        nodePointer newRoot = nod->lc;
        nod->lc = newRoot->rc;
        newRoot->rc = nod;

        update(nod);
        update(newRoot);

        return newRoot;
    }

    nodePointer r_l_rotate(nodePointer nod) {  // 右子树上插入左儿子,先(向)右旋右儿子,再左旋自己
        nod->rc = right_rotate(nod->rc);
        return left_rotate(nod);
    }

    nodePointer l_r_rotate(nodePointer nod) {  // 左子树插入右儿子
        nod->lc = left_rotate(nod->lc);
        return right_rotate(nod);
    }

    nodePointer get_new() {
        if (bintop) {
            return rubbish_bin[bintop--];
        } else {
            ++tot;
            return tot;
        }
    }

    nodePointer get_max_nod(nodePointer nod) {
        while (nod->rc != null)
            nod = nod->rc;
        return nod;
    }
    nodePointer get_min_nod(nodePointer nod) {
        while (nod->lc != null)
            nod = nod->lc;
        return nod;
    }

    nodePointer ins(nodePointer& nod, Type val) {
        if (nod == null) {
            // std::cout << "[INS] " << val << '\n';
            nod = get_new();
            nod->val = val;
            nod->siz = 1;
            nod->lc = nod->rc = null;
        } else if (val > nod->val) {
            nod->rc = ins(nod->rc, val);
            if (get_height(nod->rc) - get_height(nod->lc) == 2) {
                if (val > nod->rc->val)
                    nod = left_rotate(nod);
                else
                    nod = r_l_rotate(nod);
            }

        } else {
            nod->lc = ins(nod->lc, val);
            if (get_height(nod->lc) - get_height(nod->rc) == 2) {
                if (val < nod->lc->val)
                    nod = right_rotate(nod);
                else
                    nod = l_r_rotate(nod);
            }
        }
        update(nod);
        return nod;
    }

    nodePointer& del(nodePointer& nod, Type val) {
        if (nod == null) {
            return null;
        } else {
            if (val == nod->val) {
                if (nod->lc != null && nod->rc != null) {
                    if (get_height(nod->lc) > get_height(nod->rc)) {
                        nodePointer pre = get_max_nod(nod->lc);  // 如果左子树比右子树大,则从左子树找最大值替换当前节点
                        nod->val = pre->val;
                        nod->lc = del(nod->lc, pre->val);  // 递归删除刚才找到的最大值
                    } else {
                        nodePointer pre = get_min_nod(nod->rc);
                        nod->val = pre->val;
                        nod->rc = del(nod->rc, pre->val);
                    }
                } else {
                    // std::cout << "[DEL] " << nod->val << '\n';
                    rubbish_bin[++bintop] = nod;
                    if (nod->lc != null)
                        nod = nod->lc;
                    else if (nod->rc != null)
                        nod = nod->rc;
                    else
                        return null;
                }
            } else if (val > nod->val) {
                nod->rc = del(nod->rc, val);
                // 删除右子树节点后左子树高于右子树,相当于在左子树插入新节点,对应单右旋或左旋后右旋
                if (get_height(nod->lc) - get_height(nod->rc) > 2) {
                    if (get_height(nod->lc->rc) > get_height(nod->lc->lc))
                        nod = l_r_rotate(nod);
                    else
                        nod = right_rotate(nod);
                }
            } else {
                nod->lc = del(nod->lc, val);
                if (get_height(nod->rc) - get_height(nod->lc) > 2) {
                    if (get_height(nod->rc->lc) > get_height(nod->rc->rc))
                        nod = r_l_rotate(nod);
                    else
                        nod = left_rotate(nod);
                }
            }
            return nod;
        }
    }

    nodePointer find(nodePointer nod, Type val) {
        // if (val == nod->val || nod == null)
        //     return nod;
        // else if (val > nod->val)
        //     return find(nod->lc, val);
        // else
        //     return find(nod->rc, val);
        while (nod->val != val) {
            if (nod->val > nod)
                nod = nod->lc;
            else
                nod = nod->rc;
        }
    }

    nodePointer find_rk(nodePointer nod, int k) {
        while (nod->lc->siz != k - 1) {
            while (!nod->lc)
                nod = nod->rc;
            if (nod->lc->siz > k - 1)
                nod = nod->lc;
            else
                k -= nod->lc->siz + 1, nod = nod->rc;
        }
        return nod;
    }

    void print(nodePointer nod) {
        if (nod == null) return;
        if (nod->lc != null) std::cout << nod->val << ' ' << nod->lc->val << std::endl;
        if (nod->rc != null) std::cout << nod->val << ' ' << nod->rc->val << std::endl;
        print(nod->lc);
        print(nod->rc);
    }

public:
    AVLTree() {
        tot = tree;
        null = root = tot;

        // init null node
        null->lc = null->rc = root->lc = root->rc = null;
        null->height = null->val = null->siz = 0;
    }

    void insert(Type val) {
        ins(root, val);
    }
    void remove(Type val) {
        del(root, val);
    }
    int get_rank(Type val) {
        nodePointer nod = root;
        Type ans = 0;
        while (nod != null) {
            if (nod->val == val)
                return get_siz(nod->lc) + ans + 1;
            else if (nod->val > val)
                nod = nod->lc;
            else {
                ans += get_siz(nod->lc) + 1;
                nod = nod->rc;
            }
        }
        return ans + 1;
    }
    Type find_by_rank(int rk) {
        return find_rk(root, rk)->val;
    }
    Type get_front(Type val) {
        nodePointer nod = root;
        Type ans;
        while (nod != null) {
            if (nod->val == val)
                return get_max_nod(nod->lc)->val;
            else if (nod->val > val)
                nod = nod->lc;
            else {
                ans = nod->val;
                nod = nod->rc;
            }
        }
        return ans;
    }
    Type get_next(Type val) {
        nodePointer nod = root;
        Type ans;
        while (nod != null) {
            if (nod->val == val)
                return get_min_nod(nod->rc)->val;
            else if (nod->val < val)
                nod = nod->rc;
            else {
                ans = nod->val;
                nod = nod->lc;
            }
        }
        return ans;
    }

    void debug() {
        std::cout << get_siz(root) << '\n' << root->val << '\n';
        print(root);
    }
};

AVLTree<int, 100005> tree;
int n, m;
int opt, x;
int ans;

using std::cin;
using std::cout;

#define endl '\n'

int main() {
    // std::ios::sync_with_stdio(0);
    // cin.tie(nullptr);
    // cout.tie(nullptr);

    cin >> n;
    while (n--) {
        cin >> opt >> x;
        if (opt == 1) {
            tree.insert(x);
        } else if (opt == 2) {
            tree.remove(x);
        } else if (opt == 3) {
            cout << tree.get_rank(x) << endl;
        } else if (opt == 4) {
            cout << tree.find_by_rank(x) << endl;
        } else if (opt == 5) {
            cout << tree.get_front(x) << endl;
        } else if (opt == 6) {
            cout << tree.get_next(x) << endl;
        }
    }

    tree.debug();

    return 0;
}

2023/5/4 19:21
加载中...