treeを用いて区間内の特定の要素の個数をlogの2乗で求めるデータ構造

はじめに

こんにちは。mofuuniです。
今回は、タイトルにもある通り 「treeを用いて区間内の特定の要素の個数を$O(log^2N)$で求めるデータ構造」 というものを作ったので紹介したいと思います。

※2024/4/4追記 多数の間違いを発見したため大幅に修正しました。

できること

次のクエリに答えることができます。

1. $A_{pos}$の値を$A_{pos}⋆x$に変更せよ。
2. $A_l, A_{l+1}, ..., A_r$の中に含まれる$x$の個数を答えよ。

計算量はどちらも$O(log^2N)$です。

実装

#include <bits/stdc++.h>
#include<ext/pb_ds/assoc_container.hpp>
#include<ext/pb_ds/tree_policy.hpp>
#include<ext/pb_ds/tag_and_trait.hpp>
using namespace std;
using namespace __gnu_pbds;
template<class T> using Tree=tree<T,null_type,less<T>,rb_tree_tag,tree_order_statistics_node_update>;

template<class T>
class FreqencyCountTree{
    public:
    using OP=function<T(T,T)>;
    vector<T> vec;
    map<T,Tree<int>> data;
    OP op;
    FreqencyCountTree(int n,T init):vec(n+1,init){
        for(int i=1; i<=n; i++)data[init].insert(i);
        op=[](T a,T b){return b;};
    }
    FreqencyCountTree(int n,T init,OP op_):vec(n+1,init),op(op_){
        for(int i=1; i<=n; i++)data[init].insert(i);
    }
    FreqencyCountTree(vector<T> v):vec(v){
        for(int i=1; i<v.size(); i++)data[v[i]].insert(i);
        op=[](T a,T b){return b;};
    }
    FreqencyCountTree(vector<T> v,OP op_):vec(v),op(op_){
        for(int i=1; i<v.size(); i++)data[v[i]].insert(i);
    }
    void update(int pos,T x){
        data[vec[pos]].erase(pos);
        vec[pos]=op(vec[pos],x);
        data[vec[pos]].insert(pos);
    }
    int count(int l,int r,T x){
        if(l>r)return 0;
        if(!data.count(x))return 0;
        int lef=data[x].order_of_key(l);
        int rig=data[x].order_of_key(r+1);
        return rig-lef;
    }
};

名前は仮です。

説明

vec

元の配列のコピーです。

data

$A_{pos} = x$のとき$data[x]$の中に$pos$を入れます。
つまり、「$A$の何番目が$x$か」ということです。

op

値を変更する演算です。指定しなければ代入になります。

update関数

クエリ1をこなします。計算量は$O(log^2N)$です。

count関数

クエリ2をこなします。
treeのorder_of_key関数により$data[x]$内の$l$未満の要素の個数($l$より前に$x$が何個あるか)、$r$以下の要素の個数($r$までに$x$が何個あるか)を数え、引くことで区間内の$x$の個数が求まります。
計算量は$O(log^2N)$です。

verify(問題のネタバレを含みます)

verify

ABC248 D - Range Count Query
要素の個数のみです。

ABC343 F - Second Largest Query
想定解ではSegment Treeに4つも情報を載せるので大変なところを楽できます。
ちなみにこのデータ構造はこの問題を考えているときに思いつきました。

C: Range Same Query

ABC157 E - Simple String Query
数列以外でも使えます。

おわりに

初めての執筆でわかりにくいところも多々あったかと思いますが、最後まで読んでくださりありがとうございました。
ご指摘・ご質問等ございましたらお気軽にどうぞ。