一种无脑的平衡树做法

stripe_python · 2023-08-12 20:26:08 · 题解

Updated 2023/8/14：因为没介绍平板电视没过审，提交后发现题解通道关了，不过代码是能过 Hack 的。

Updated 2023/8/22：修改代码。

Step1：代码结构分析

看到题是典型的数据结构模拟，让我们支持 start、push、pop 三种操作，均摊 O(\log{n}) 及以下。

这种模拟题本蒟蒻的习惯是面向对象写，先设计一种类，如下所示：

class Manager {
public:
    pair<string, string> start();
    bool push(const string& s);
    bool pop(const string& s);
};

start 函数返回两个上场人员的名字，可以返回空字符串；
push 函数将人员加入队列，返回是否成功；
pop 函数将人员移出队列，返回是否成功；

于是我们的代码雏形如下：

#include <bits/stdc++.h>
#define endl '\n'
// endl定义为换行符防止刷新缓冲，卡常小技巧 
using namespace std;

class Manager {
public:
    pair<string, string> start();
    bool push(const string& s);
    bool pop(const string& s);
};

int n;
string c, opt;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(0), cout.tie(0);
    cin >> n;
    Manager man;
    // 这一块就是纯调API了，不解释
    while (n--) {
        cin >> opt;
        if (opt == "start") {
            auto h = man.start();
            if (h.first.empty()) {
                cout << "Error" << endl;
            } else if (h.second.empty()) { 
                cout << h.first << endl;
            } else {
                cout << h.first << ' ' << h.second << endl;
            }
        } else if (opt == "arrive") {
            cin >> c;
            bool ok = man.push(c);
            if (ok) cout << "OK" << endl;
            else cout << "Error" << endl;
        } else if (opt == "leave") {
            cin >> c;
            bool ok = man.pop(c);
            if (ok) cout << "OK" << endl;
            else cout << "Error" << endl;
        }
    }
    return 0;
}

Step2：数据结构分析

考虑用二元组 (x, user) 记录一个人，x 记录入队时间，user 记录人名。

定义结构体如下：

struct node {
    int x;
    string user;

    node(int _x, string _u) : x(_x), user(_u) {}

    bool operator< (const node& other) const {
        return x < other.x;   // 这里按入队时间排序
    }
};

现在所有要支持的操作摆出来了：

现在来具体分析一下各个操作：

push：获取队列中的最后一个 x，并将 (x + 1, user) 放入队列；
pop：根据哈希表获取 s 的 x，删除 (x, s)；
start：设 x, y 表示当前在玩的玩家，先将 x, y 入队，接下来取出队列的前两项，pop出来返回即可。

考虑使用哈希表记录每个 user 的 x，又要记录排队顺序，还要 O(\log{n}) 查找，首先想到set。

~~然而本蒟蒻太弱了，比赛时调两个小时的set调不出来QwQ~~

由上，支持对数的增删改查——平衡树！

Step3：平衡树实现

注意到，本题元素不重复——
平板电视大法好！

#include <bits/extc++.h>
using namespace __gnu_pbds;
#define treetype tree<node, null_type, less<node>, rb_tree_tag, tree_order_statistics_node_update>

简单介绍一下：

null_type：无映射；
rb_tree_tag：红黑树；
tree_order_statistics_node_update：用这个才支持查询排名。

再介绍一下 tree 的常用操作：

erase：从树中移出给定的 node；
insert：插入给定的 node；
find_by_order：根据排名找 node，注意返回的是指针；
order_by_key：根据 node 获取排名。

tree 的常数略大，但本题来说可以接受 O(n\log{n}) 的复杂度。

详细介绍见 OI Wiki。

Step4：AC Code

思路讲过了，代码就不放太多注释了，自认为码风还是可读性较高的。

只放 Manager 的代码。

using namespace __gnu_pbds;   // 平板电视大法好！ 

class Manager {
private:
    struct node {
        int x;
        string user;

        node(int _x, string _u) : x(_x), user(_u) {}

        bool operator< (const node& other) const {
            return x < other.x;    // 这里按入队时间排序
        }
    };
    string x, y;
#define treetype tree<node, null_type, less<node>, rb_tree_tag, tree_order_statistics_node_update>
    treetype tr;    // 平衡树 
    map<string, int> st;   // 记录入队时间 

    void show() {   // 调试用的 
        for (int i = 0; i < tr.size(); i++) {
            cout << tr.find_by_order(i)->user << " ";
        }
        cout << endl;
    }

    bool find(const string& s) {
        return st.count(s);
    }

    bool _pop(const string& s) {   
        if (tr.empty() || s.empty()) return false;
        if (!find(s)) return false;
        tr.erase(node(st[s], s));    // 删除(st[s], s) 
        st.erase(st.find(s));    // 哈希表一并要删除 
        return true;
    }

    bool _push(const string& s) {
        if (find(s)) return false;
        if (s.empty()) return false;
        //show();
        // 计算当前入队时间 
        int w = tr.empty() ? 0 : (tr.find_by_order(tr.size() - 1)->x) + 1;
        tr.insert(node(w, s));   // 插入(w, s) 
        st[s] = w;   // 设置哈希表 
        return true;
    }

public:
    pair<string, string> start() {
        _push(x), _push(y);
        //show();
        if (tr.empty()) return {"", ""};
        else if (tr.size() == 1) x = tr.find_by_order(0)->user, _pop(x);
        else x = tr.find_by_order(0)->user, y = tr.find_by_order(1)->user, _pop(x), _pop(y);
        return {x, y};
    }

    bool push(const string& s) {
        if (x == s || y == s) return false;
        return _push(s);
    }

    bool pop(const string& s) {
        if (x == s || y == s) return false;
        return _pop(s);
    }
};