Leetcode LCP 52. 二叉搜索树染色 珂朵莉树模板题

能比较高效的解决一下问题

• 区间加
• 区间赋值
• 求区间第k大值

typedef long long ll;
struct ChthollyTree{
    struct node
    {
        ll l, r;
        mutable ll v;
        node(ll l, ll r, ll v) : l(l), r(r), v(v) {}
        bool operator<(const node &o) const { return l < o.l; }
    };
    set<node> tree;
    void insert(ll l,ll r,ll v){
        tree.insert({l,r,v});
    }
    auto split(ll pos)
    {
        auto it = tree.lower_bound(node(pos, 0, 0));
        if (it != tree.end() && it->l == pos)
            return it;
        it--;
        ll l = it->l, r = it->r, v = it->v;
        tree.erase(it);
        tree.insert(node(l, pos - 1, v));
        return tree.insert(node(pos, r, v)).first;
    }
    ll get_sum(ll l, ll r)
    {
        long long ans=0;
        auto end = split(r + 1), begin = split(l),it = begin;
        for(;it!=end;it++)
            ans+=it->v*(it->r-it->l+1);
        return ans;
    }
    void assign(ll l, ll r, ll v)
    {
        auto end = split(r + 1), begin = split(l);
        tree.erase(begin, end);
        tree.insert(node(l, r, v));
    }
    void add(ll l, ll r, ll v)
    {
        auto end = split(r + 1);
        for (auto it = split(l); it != end; it++)
            it->v += v;
    }
    ll kth(ll l, ll r, ll k)
    {
        auto end = split(r + 1);
        vector<pair<ll, ll>> v;
        for (auto it = split(l); it != end; it++)
            v.push_back(make_pair(it->v, it->r - it->l + 1));
        sort(v.begin(), v.end());
        for (int i = 0; i < v.size(); i++)
        {
            k -= v[i].second;
            if (k <= 0)
                return v[i].first;
        }
        return 0;
    }
};
class Solution {
public:
    vector<int> a;
    unordered_map<int,int> mp;
    void dfs(TreeNode *p){
        if(p==NULL) return;
        dfs(p->left);
        a.push_back(p->val);
        dfs(p->right);
    }
    int getNumber(TreeNode* root, vector<vector<int>>& ops) {
        dfs(root);
        int n=a.size();
        for(int i=0;i<n;i++) mp[a[i]]=i;
        ChthollyTree tree;
        tree.insert(0,n-1,0);
        for(auto op:ops){
            int t=op[0],l=mp[op[1]],r=mp[op[2]];
            tree.assign(l,r,t);
        }
        return tree.get_sum(0,n-1);
    }
};