問題概要

$ n $要素からなる数列$ a = \{ a_1, a_2, \cdots, a_n \} $が与えられる．$ a_i $は，テレビシリーズのシーズン$ i $のエピソードが第$ a_i $エピソードまで存在することを表す．

シーズン$ i $のエピソード$ j $とシーズン$ j $のエピソード$ i $が両方存在するようなペアの数を求める．

制約

$ 1 \leq n \leq 2 \times 10^5 $

$ 1 \leq a_i \leq 10^9 $

解法概要

まず，$ a_i $ が$ n $を超える場合は$ n $と考えてよい．シーズンは $ n $までしかないのだから，シーズン$ i$のエピソード$ j (> n) $とペアになるシーズン・エピソードの組は存在しない．以下では$ a_i \leq n $として解説を行う．

題意を満たすペアの条件から，ペアの数は以下のような式で表現できる．

$ \sum_{j=1}^n \sum_{k=j+1}^{a_j} (a_k \geq j ? 1 : 0)$

(丸括弧内の?と:の意味はC言語の三項演算子と同じ意味だと思ってください)

しかしこれをこのまま実装しようとしようとすると$ \mathcal{O} (n^2) $となり，$ 1 \leq n \leq 2 \times 10^5 $であるこの問題では間に合わない．そこで，条件を満たす要素の数え上げにBIT (Binary Indexed Tree, Fenwick Tree) を使うことを考える(BITそのものに関する解説は他サイトさんに譲ります)．結論から言えば，BITを使うことで2つ目のシグマに当たる部分が$ \mathcal{O} (\log n) $で計算できるようになるので，全体で$ \mathcal{O} (n \log n) $で解が求められる．

$ \sum_{k=j+1}^{a_j} (a_k \geq j ? 1 : 0)$の部分については，$ j+1 \leq k \leq a_j $である$ k $の中で$ a_k \geq j $となるような$ k $がいくつあるかを数えればよい．これを実現するためには，BITに要素を追加する順番や和を求める順番を工夫しなくてはならない．

$ a_k \geq j $となる$ k $を数えたいのだから，$ a_k $の値の大きいものからBITに追加したい．そのための下準備として，インデックスから値がわかる配列$ a $以外にもある値を持つインデックスがわかる2次元配列$ b $を用意しておく．また，対応するペアの数はシーズン$ n $からシーズン$ 1 $へ，という順番で求める．

シーズン$ i $のエピソードのどれかとペアになるエピソードの集合を求めるとき，最初に$ a_k = j$となるような$ k $について，BITの$ k $番目の要素に1をを加えておき，BITに足された値の$ j+1 $番目から$ a_j $番目まで (境界を含む) の和を求めればよい．

ソースコード

#include <bits/stdc++.h>

#define rep(n) for(int i=0;i<n;i++)
#define repp(j, n) for(int j=0;j<n;j++)
#define reppp(i, m, n) for(int i=m;i<n;i++)
#define all(c) c.begin(), c.end()
#define rall(c) c.rbegin(), c.rend()
#define debug(x) cerr << #x << ": " << x << endl;

using namespace std;

typedef long long ll;
typedef pair<ll, ll> Pll;
typedef pair<int, int> Pii;

struct BIT{
    int n;
    int bit[1000001];
    BIT(int n){
        this->n = n;
        fill(bit, bit+n, 0);
    }

    void add(int idx, int x){
        for(int i=idx;i<=this->n;i+=i&-i) bit[i] += x;
    }

    //bit[1]からbit[end]までの和 (閉区間)
    int sum(int end){
        int ret = 0;
        for(int i=end;i>=1;i-=i&-i) ret += bit[i];
        return ret;
    }

    int count_swap(vector<int> v){
        int ret = 0;
        reppp(i, 1, n){
            add(v[i], 1);
            ret += i - sum(v[i]);
        }
        return ret;
    }
};

int main(){
    int n;
    cin >> n;
    int a[n+1];
    vector<int> b[n+1];
    reppp(i, 1, n+1){
        cin >> a[i];
        a[i] = min(n, a[i]);
        b[a[i]].push_back(i);
    }

    ll ans = 0;
    BIT tree = BIT(1<<18);
    for(int i=n;i>=1;i--){
        for(int bi: b[i]){
            tree.add(bi, 1);
        }

        if(i+1 <= a[i]) ans += tree.sum(a[i]) - tree.sum(i);
    }

    cout << ans << endl;
}

感想

BITのあまり自明でない利用のいい練習になった．あと式変形が楽しかった．

NoiminのNoise

Educational Codeforces Round 41 E. Tufurama

問題概要

解法概要

ソースコード

感想