「連続する整数」が多い集合を管理するデータ構造! 【備忘録】

まえがき

 AtCoder ABC194のE-"Mex Min"を解いてたら、次のお方の記事で「ある整数の集合におけるmex(minimum excluded:集合に含まれない最小の整数)を対数時間で求められるデータ構造」を知りました。

rsk0315.hatenablog.com


 他の方のコードも参考にしながら自分なりに作ってみたので、備忘録も兼ねて残しておこうと思います。



コード

 アイデアは先のブログさまのものと一緒ですのでそちらをご参考ください......!
(※考えてる範囲は閉区間です!)

/*連続する整数が多い集合を管理するデータ構造.*/
class BigSet{
    std::set<std::pair<int,int> > s;//s:整数の集合, 連続する範囲[l,r]をstd::pairで管理する.
    int inf=(1<<30);//番兵用.
    
public:
    //constructor.
    BigSet(){
        s.emplace(-inf,-inf);
        s.emplace(inf,inf);
    }
    
    bool insert(int x){//挿入. O(logN).
        auto nitr=s.lower_bound((std::make_pair(x+1,x+1)));
        auto itr=std::prev(nitr);
        auto [l,r]=*itr;
        auto [nl,nr]=*nitr;
        if(l<=x and x<=r) return false;//既に集合に含まれている場合.
        if(r==x-1){
            if(nl==x+1){
                s.erase(itr);
                s.erase(nitr);
                s.emplace(l,nr);
            }else{
                s.erase(itr);
                s.emplace(l,x);
            }
        }else{
            if(nl==x+1){
                s.erase(nitr);
                s.emplace(x,nr);
            }else{
                s.emplace(x,x);
            }
        }
        return true;
    }
    bool insert(int l,int r){//範囲[l,r]の整数を挿入. O(logN).
        assert(l<=r);
        if(l==r) return insert(l);
        auto itr1=std::prev(s.lower_bound((std::make_pair(l+1,l+1))));
        auto [l1,r1]=*itr1;
        if(l1<=l and r<=r1) return false;//完全に集合に含まれている場合.
        auto itr3=s.lower_bound((std::make_pair(r+1,r+1)));
        auto itr2=std::prev(itr3);
        auto [l2,r2]=*itr2;
        auto [l3,r3]=*itr3;
        if(l1<=l and l<=r1+1){
            if(l3==r+1){
                itr3++;
                s.erase(itr1,itr3);
                s.emplace(l1,r3);
            }else{
                s.erase(itr1,itr3);
                if(l2<=r and r<=r2) s.emplace(l1,r2);
                else s.emplace(l1,r);
            }
        }else{
            itr1++;
            if(l3==r+1){
                itr3++;
                s.erase(itr1,itr3);
                s.emplace(l,r3);
            }else{
                s.erase(itr1,itr3);
                if(l2<=r and r<=r2) s.emplace(l,r2);
                else s.emplace(l,r);
            }
        }
        return true;
    }
    bool erase(int x){//削除. O(logN).
        auto itr=std::prev(s.lower_bound((std::make_pair(x+1,x+1))));
        auto [l,r]=*itr;
        if(x<l or r<x) return false;//集合に含まれていない場合.
        s.erase(itr);
        if(l!=x) s.emplace(l,x-1);
        if(r!=x) s.emplace(x+1,r);
        return true;
    }
    bool erase(int l,int r){//範囲[l,r]の整数を削除. O(logN).
        assert(l<=r);
        if(l==r) return erase(l);
        auto itr1=std::prev(s.lower_bound((std::make_pair(l+1,l+1))));
        auto itr3=s.lower_bound((std::make_pair(r+1,r+1)));
        auto itr2=std::prev(itr3);
        auto [l1,r1]=*itr1;
        auto [l2,r2]=*itr2;
        if(itr1==itr2 and r1<l) return false;//集合に完全に含まれていない場合.
        if(l1<=l and l<=r1){
            s.erase(itr1,itr3);
            if(l1<l) s.emplace(l1,l-1);
        }else{
            itr1++;
            s.erase(itr1,itr3);
        }
        if(l2<=r and r<r2) s.emplace(r+1,r2);
        return true;
    }
    bool contains(int x) const{//集合に整数xが含まれるか判定する. O(logN).
        auto [l,r]=*std::prev(s.lower_bound((std::make_pair(x+1,x+1))));
        return (l<=x and x<=r);
    }
    int mex(int x=0) const{//Minimum Excluded. x以上であって,集合に含まれない最小の整数を求める. O(logN).
        auto [l,r]=*std::prev(s.lower_bound((std::make_pair(x+1,x+1))));
        if(l<=x and x<=r) return r+1;
        else return x;
    }
    void print() const{
        for(const auto &[l,r]:s) fprintf(stderr,"[%d,%d]",l,r);
        std::cerr<<std::endl;
    }
};
あとがき

 一応確認しましたがバグってたら教えてくださると有難いです。