部分列のMEXをセグ木で計算する方法 (Codeforces Round #678 (Div. 2) E. Complicated Computations)

問題

codeforces.com

長さ N の配列 A が与えられる。
 A の空でない任意の部分列のMEXからなる数列のMEXを計算してください。

解法

RMQのセグ木を使ってなんやかんやええかんじにする。

  • 数字をkey, index を value としてセグ木に乗せる(この表現でいいのかわかんない)

  • 配列を前から見ていく。

  • 各値について、最後に出てきた位置として-1で初期化しておく。

  • 部分列のMEXで  a_i を作れるかどうかは、前回出てきた  a_iの位置とiの間に  0から a_i-1が全て存在すればよい。

  • a_iが最後に出てきた位置を更新する。

  • 一番最後に、一番最後に出てきたiの位置からn-1の位置まででMEXをiにできるかを調べる

コード

#include<bits/stdc++.h> 
using namespace std;
typedef long long ll;

template < typename Monoid >
struct Segtree{
    using FM = function < Monoid(Monoid, Monoid) >;
    using FA = function < void(Monoid&, Monoid) >;
    int N;
    vector<Monoid> dat;
    
    const FM fm;
    const FA fa;
    const Monoid M1;

    Segtree(int n,const FM fm, const FA fa, const Monoid &M1): fm(fm), fa(fa), M1(M1){
        N = 1;
        while(N < n) N *= 2;
        dat.assign(2*N,M1);
    }

    void set(int idx, const Monoid &x){
        dat[idx+N-1] = x;
    }

    void build(){
        for(int i=N-2; i>=0; i--){
            dat[i] = fm(dat[2*i+1], dat[2*i+2]);
        }
    }

    void update(int idx, const Monoid &x){
        idx += N-1;
        fa(dat[idx],x);
        while(idx > 0){
            idx = (idx-1)/2;
            dat[idx] = fm(dat[idx*2 + 1], dat[idx*2+2]);
        }
    }

    Monoid get(int a, int b, int k, int l, int r){
        if(r<=a || b<=l) return M1;
        if(a<=l && r<=b) return dat[k];
        else{
            Monoid vl = get(a,b,k*2+1,l,(l+r)/2);
            Monoid vr = get(a,b,k*2+2,(l+r)/2,r);
            return fm(vl,vr);
        }
    }
    Monoid get(int a,int b){ return get(a,b,0,0,N); }

    Monoid at(int k) {return dat[k+N-1];}

    void debug(){
        for(int i=N-1;i<dat.size();i++){
            cerr << dat[i] << ", ";
        }
        cerr << "\n";
    }
};

// Range Minimum Query
auto fm = [](ll a,ll b){ return min(a,b); };
auto fa = [](ll& a, ll b){ a=b; };
Segtree<ll> seg(100010, fm, fa, 1e9);


signed main(){
    ios::sync_with_stdio(false);
    cin.tie(0);
    cout << fixed << setprecision(20);
    
    int n;
    cin>>n;
    int a[n];
    bool one = true;
    for(int i=0;i<n;i++){
        cin>>a[i];
        a[i]--;
        if(a[i] != 0)one = false;
    }
    bool ans[n+5]={};
    for(int i=0;i<=n+5;i++)seg.set(i,-1);
    seg.build();

    // 0 以外があるならばMEXを0にできる
    if(one == false)ans[0] = 1;
    for(int i=0;i<n;i++){
        // 0 は、やばい
        if(a[i] == 0){
            seg.update(0,i);
            continue;
        }
        // 0, ..., a[i]-1 で、最も左にあるもの(出ていないなら-1)の位置を調べる
        ll mi = seg.get(0,a[i]);

        // 最後に出てきたa[i]の位置を調べる
        ll pre = seg.at(a[i]);

        // pre より右側でだけで 0, ..., a[i]-1を全て持っているなら、MEXはa[i]
        if(pre < mi)ans[a[i]] = 1;

        seg.update(a[i],i);
    }
    
    // 1,...,n+1 の全てで、一番最後に出てきた位置からn-1の位置までで、MEXをiにできるか
    // 0 は、やばい
    for(int i=1;i<=n+1;i++){
        ll mi = seg.get(0,i);
        ll pre = seg.at(i);
        if(pre < mi)ans[i] = 1;
    }

    // 最大でn+1 (0-indexed)
    for(int i=0;i<=n+5;i++){
        if(ans[i] == 0){
            cout << i+1 << endl;
            return 0;
        }
    }
}