C++勢だってsetのk番目を取得したい！

「実行速度がとにかく速い」
「1番使っている人が多いから解説も多い」
「雑なコードでも時間制限に間に合うことがある」
競プロで有利と名高く、私もずっと愛用してきたC++。
Pythonを使う人も多いですが、実行速度の安心感は大きいです。
しかしそれでも

「Python勢が羨ましい」

と思うタイミングがあります。それは、順序付き集合に対して

s[k]

のようなアクセスをしたいときです。

C++の std::set や std::multiset は順序を保ってくれますが、$k$番目の要素を直接取得することはできません。

無理やり取得しようとするとイテレータを1つずつ進める必要があり、最悪で集合サイズを $N$ として $O(N)$ の時間がかかってしまいます。

どうやらPythonには標準ではないもののAtCoderのジャッジ環境で広く使われているSortedSetという便利なものがあるらしく。
羨ましかったので、C++でも作ることにしました。

この記事では、競プロで役に立つ「インデックスアクセスができる set / multiset」について紹介します。

完成したもの

#pragma GCC optimize("O3,unroll-loops")
#pragma GCC target("avx2")
#include<bits/extc++.h>
using namespace std; using namespace __gnu_pbds;
#define T tree<conditional_t<U,K,pair<K,int>>,null_type,conditional_t<U,C,less<>>,rb_tree_tag,tree_order_statistics_node_update>
#define g(k,...) ([&]{if constexpr(U)return k;else return pair{k,0 __VA_ARGS__};}())
#define key(it) (*(K*)&*it)
#define K_ b::order_of_key
template<class K,bool U,class C=less<>>
struct B:T{
    using b=T;using b::b;int id{};
    b* mut()const{return(b*)this;}
    struct I{
        using iterator_category = random_access_iterator_tag;
        using value_type = K; using difference_type = int;
        using pointer = const K*; using reference = const K&;
        const B* p; typename b::iterator it; int r;
        auto& operator*()const{return key(it);}
        auto operator->()const{return &key(it);}
        I operator+(int n)const{int x=r+n;return{p,(unsigned)x>=p->size()?p->mut()->end():p->mut()->find_by_order(x),x};}
        I operator-(int n)const{return*this+-n;}
        I& operator+=(int n){return*this=*this+n;}
        I& operator-=(int n){return*this=*this+-n;}
        friend I operator+(int n,I i){return i+n;}
        int operator-(I o)const{return r-o.r;}
        I& operator++(){return*this+=1;}
        I& operator--(){return*this+=-1;}
        I operator++(int){auto tmp=*this;*this+=1;return tmp;}
        I operator--(int){auto tmp=*this;*this+=-1;return tmp;}
        auto operator<=>(const I& o)const{return r<=>o.r;}
        bool operator==(const I& o)const{return r==o.r;}
        auto& operator[](int n)const{return*(*this+n);}
        operator typename b::iterator()const{return it;}
    };
    using iterator = I; using const_iterator = I;
    using reverse_iterator = std::reverse_iterator<I>; using const_reverse_iterator = reverse_iterator;
    I begin(int=0)const{return{this,mut()->begin(),0};} I end(int=0)const{return{this,mut()->end(),(int)b::size()};}
    I cbegin()const{return begin();} I cend()const{return end();}
    reverse_iterator rbegin()const{return reverse_iterator(end());} reverse_iterator rend()const{return reverse_iterator(begin());}
    reverse_iterator crbegin()const{return rbegin();} reverse_iterator crend()const{return rend();}
    I make_I(typename b::iterator it)const{return{this,it,it==mut()->end()?(int)b::size():(int)K_(*it)};}
    auto& operator[](int i)const{return key(mut()->find_by_order(i));}
    auto insert(auto&& x){return b::insert(g(x,+id++));}
    I find(auto&& k)const{auto it=lower_bound(k);return it!=end()&&*it==k?it:end();}
    I lower_bound(auto&& k)const{return make_I(mut()->lower_bound(g(k)));}
    I upper_bound(auto&& k)const{if constexpr(U)return make_I(mut()->upper_bound(k));else return make_I(mut()->lower_bound(g(k,+2e9)));}
    int count(auto&& k)const{if constexpr(U)return b::count(k);else return K_(g(k,+2e9))-K_(g(k));}
    int erase(auto&& k){if constexpr(U)return b::erase(k);int l=lower_bound(k).r,r=upper_bound(k).r,c=r-l;for(int i=0;i<c;++i)b::erase(mut()->find_by_order(l));return c;}
    void erase(I it) { b::erase((typename b::iterator)it); }
};
template<class K,class C=less<>>using _s=B<K,1,C>;
template<class K,class C=less<>>using _m=B<K,0,C>;
#undef T
#undef g
#undef key
#undef K_
#define set _s
#define multiset _m

これをコード冒頭に貼り付けることでset,multisetにoperator[]でアクセスできます。
※GNU拡張を使用しているためGCCでのコンパイルが必須になります。

仕組み

GNU PBDSのOrder Statistic Treeを利用しています。
Order Statistic Treeはstd::setとは違い内部で

各ノードに部分木サイズを持たせた赤黒木

を使用しています。
そのため「$k$番目の要素」を部分木サイズを見ながら木を下るだけで求められます。
集合サイズを$N$とすると計算量は木の高さ$log N$に依存するので結果として
$O(log N)$
で求められるという仕組みです。

実際に使える問題

ここからは実際にテンプレートが役立つ問題を紹介します。

ABC 458 D - Chalkboard Median

問題文

黒板に整数 $X$ が 1 つ書かれています。
$Q$ 個のクエリが与えられるので、順に処理してください。
$i$ 個目 $(1 \le i \le Q)$ のクエリは以下の通りです。

2 つの整数 $A_i, B_i$ が与えられる。黒板に新たに 2 つの整数 $A_i, B_i$ を書く。
その後、黒板に書かれた $2i+1$ 個の整数の中央値を出力する。

制約

• $1 \le X \le 10^9$
• $1 \le Q \le 2 \times 10^5$
• $1 \le A_i, B_i \le 10^9$
• 入力される値は全て整数

この問題の解法は

• 優先度付きキューを使用した解法
• 座標圧縮をしてFenwick Tree上で二分探索する解法

などが挙げられます。
ですが、インデックスアクセス可能な順序付き多重集合でも解けます。
具体的には2つの整数を書くときにinsertを使用し、中央値はインデックスでアクセスするだけです。

// ここに先程のテンプレートを貼り付ける
int main(){
    int X, Q;
    cin >> X >> Q;
    multiset<int> st;
    st.insert(X);
    for (int i = 0; i < Q; i++) {
        int A, B;
        cin >> A >> B;
        st.insert(A); st.insert(B);
        cout << st[i+1] << endl; // 中央値であるi+1にアクセス
    }
    return 0;
}

ABC 459 C - Drop Blocks

問題文
$N$ 個のマスが左右一列に並んでいます。はじめ、すべてのマスには何も置かれていません。
$Q$ 個のクエリが与えられるので、順に処理してください。各クエリは次の 2 種類のいずれかです。

• 1 x : 左から $x$ 番目のマスにブロックを 1 個積む。その後、すべてのマスに 1 個以上のブロックが積まれているならば、すべてのマスからブロックを 1 個ずつ取り除く。
• 2 y : $y$ 個以上のブロックが積まれているマスの個数を出力する。

制約

• $1 \le N \le 3 \times 10^5$
• $1 \le Q \le 3 \times 10^5$
• $1 \le x \le N$
• $1 \le y \le 3 \times 10^5$
• 入力はすべて整数
• 2 種類目のクエリが少なくとも 1 つ存在する

この問題は公式解説の解法を思いつく人が少なく、多くの人がBinary Indexed Treeで解いたであろう問題です。

この問題はテンプレートに実装されていた、ランダムイテレータ風イテレータを使うことで解くことができます。なぜなら、$y$個以上を高速に取得できるためです。

// ここに先程のテンプレートを貼り付ける
int main() {
    ll N, Q;
    cin >> N >> Q;
    vector<ll> block(N);
    multiset<ll> st;
    rep(i,N) st.insert(0);
    ll cnt = 0;
    rep(i,Q){
        ll type;
        cin >> type;
        if(type == 1){
            ll x;
            cin >> x; x--; // 0-indexedに合わせる
            // ブロックを1つ積む
            auto it = st.find(block[x]);
            st.erase(it);
            block[x]++;
            st.insert(block[x]);
            // 毎回段を取り除くとTLEするのでcntで取り除いた数を数える
            it = st.upper_bound(cnt);
            if(it == st.begin()) cnt++;
        }else{
            ll y;
            cin >> y;
            y += cnt; // 取り除いた分を合わせる
            auto it = st.lower_bound(y);
            // 最後からの距離が答えになる
            // 距離の取得はランダムイテレータ風の実装によってO(log N)
            // ここでもしSTLのmultisetとdistanceを使うと最悪O(N)でTLEする
            ll ans = st.end()-it;
            cout << ans << endl;
        }
    }
    return 0;
}

使用する際の注意点

インデックスアクセスできるset,multisetは普通のSTLの完全上位互換に見えますが、やはりデメリットもあります。特に気をつけてほしいのが、

部分木サイズを持つ分、定数倍やメモリ効率の悪化が起きる

ということです。
特に定数倍については気をつけなければならず、数倍遅くなるケースもあります。
そのため、必要なときだけテンプレートを貼り付けることをおすすめします。

また、GNU拡張のため、競技プログラミングではあまり問題になりませんがGCCでのコンパイルに依存することにも注意が必要です。

さらに高速に動かしたい場合は、赤黒木を捨てて

Adaptive Radix Tree

というデータ構造で1から自分で作るという方法があります。
この記事では解説しませんが気になる方は調べてみてください。

最後に

データ構造にはトレードオフがつきものです。
何でもできる万能なデータ構造は存在しません。
そのため、データ構造は正しく使い分けましょう。