左端から見た右端、右端から見た左端にそれぞれ条件があるような区間の数え上げテクニック

イマイチ良い言葉が見つからなかったのでタイトルが長くなりました。最近見る機会が多かったので記事を書いてみます。

以下のような形に帰着されるような区間の数え上げに使えるテクニックです。

例題

以下の2つの数列が与えられる。

左端から見た右端の条件を示す数列： $A_{1}, ..., A_{N}$
右端から見た左端の条件を示す数列： $B_{1}, ..., B_{N}$

区間 $\lbrack l, r \rbrack$ であって、以下の条件を満たすものの個数を求めよ。

$1 \le l \le r \le N$ である。
左端 $l$ から見た右端 $r$ の条件として、 $r \le A_{l}$ を満たす。
右端 $r$ から見た左端 $l$ の条件として、 $B_{r} \le l$ を満たす。

（制約）

$1 \le N \le 200000$ とか（ $O(N^{2})$ が間に合わないくらい）
$i \le A_{i} \le N$
$1 \le B_{i} \le i$

平たく言うと、左端に「対となる右端は○○以下でなければいけない」、右端に「対となる左端は○○以上でなければいけない」という条件がそれぞれあって、どちらも満たすような区間を数え上げるというものです。不等号の向きは逆でも良いです。

f:id:betrue12:20200328011812p:plain

解き方

BITやセグメント木など、1点加算と区間和取得ができるデータ構造を使います。

要素のインデックスを $1, 2, ..., n$ と見ていきながら、以下の処理をしていきます。

BITを用いて、まだ生きている左端のインデックスを管理する。
今見ている要素を右端とした時に採用可能な左端の個数をBITから取得し、答えに足し込む。

途中時点でのBITの状態は、例えば以下のようになります。今までに見てきた左端候補のうち「まだ生きている」もの、つまり $i \le A_{l}$ であるものだけに $1$ が立っています。 $A_{3}$ はもう死んでしまっているのでインデックス $3$ の値は $0$ になっています。

f:id:betrue12:20200328024421p:plain

これは要素を順番に見ていきながら、各 $l$ について

$l$ を見始めた時に、 $l$ 番目の位置に $1$ を加算する。
$A_{l}$ を通り過ぎた時に、それを打ち消すために $-1$ を加算する。

という処理をしていけば良いです。後者の処理を行うために、予め $A_{l}$ の値ごとに $l$ を分類しておくと良いでしょう。

そして今見ている要素を右端 $r$ とした時に採用可能な左端の個数は、BITから $\lbrack B_{r}, r\rbrack$ の区間和を取得することで求められます。

全体計算量は $O(N\log N)$ になります。

実装例

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

template<typename T>
struct BIT {
    int n;
    vector<T> dat;

    BIT(int n=0){
        initialize(n);
    }

    void initialize(int nin){
        n = nin;
        dat.resize(n, 0);
    }

    T sum(int i){
        T s = 0;
        while(i >= 0){
            s += dat[i];
            i = (i & (i+1)) - 1;
        }
        return s;
    }

    T sum_between(int i, int j){
        return sum(j) - sum(i-1);
    }

    void plus(int i, T x){
        while(i < n){
            dat[i] += x;
            i |= i+1;
        }
    }
};

int main(){
    int N;
    cin >> N;
    vector<int> A(N), B(N);
    for(int i=0; i<N; i++) cin >> A[i], A[i]--;
    for(int i=0; i<N; i++) cin >> B[i], B[i]--;

    int64_t ans = 0;
    BIT<int> bit(N);
    
    // 「ここを出る時にBITから1減らす」というインデックス
    vector<vector<int>> out(N);

    for(int i=0; i<N; i++){
        // iを左端としてBITに追加
        bit.plus(i, 1);
        out[A[i]].push_back(i);

        // iを右端として区間の個数を取得（引数が閉区間なので注意）
        ans += bit.sum_between(B[i], i);

        // 終了した左端をBITから削除
        for(int a : out[i]) bit.plus(a, -1);
    }

    cout << ans << endl;
    return 0;
}

コンテスト等での問題（ネタバレがあります）

ほとんどの場合、実際のコンテスト等での問題では、まず与えられた条件を考察して

左端 $l$ から見た右端 $r$ の条件として、 $r \le A_{l}$ を満たす。
右端 $r$ から見た左端 $l$ の条件として、 $B_{r} \le l$ を満たす。

という条件に落とし込むことが必要です。いくつか問題例を紹介します。

yukicoder No.1031 いたずら好きなお姉ちゃん

No.1031 いたずら好きなお姉ちゃん - yukicoder

スワップされる最小値と最大値の組を区間として考えます。以下の条件を満たす区間を数えれば良いことになります。

区間 $\lbrack l, r\rbrack$ の中で、 $P_{l}$ が最小値を、 $P_{r}$ が最大値を取る。
区間 $\lbrack l, r\rbrack$ の中で、 $P_{l}$ が最大値を、 $P_{r}$ が最小値を取る。

1番のケースだけを数えることにします。（2番のケースは、数列を反転させて同じ処理を繰り返せば良いです）

このとき左端 $l$ から見た $r$ の条件は「 $P_{l}$ より小さい要素が区間 $\lbrack l, r\rbrack$ に存在しないこと」なので、次に登場する $P_{l}$ より小さい要素の手前が $A_{l}$ に相当します。右端 $r$ から見た左端 $l$ の条件 $B_{r}$ も同様に考えることができます。

この $A_{l}$ や $B_{r}$ はセグメント木や優先度付きキューなどを用いることで計算できるので、この記事で説明した形式に落とし込むことができて、全体計算量 $O(N\log N)$ などで解くことができます。

（HackerRank）Thanks Kosen 2020 G - 答辞

Programming Problems and Competitions :: HackerRank

先ほどの問題設定とは、片方の不等号の向きが逆になります。

$A_{i}$ のほうは区間和が関わってくるので、累積和を取って $S_{i} = A_{1} + \cdots + A_{i-1}$ としておきます。すると満たすべき条件は

$B_{l} \le S_{r+1} - S_{l} \le B_{r}$

となるので、左右の不等号それぞれで式変形すると、

左端から見た右端の条件は、 $B_{l} + S_{l} \le S_{r+1}$ であること。
右端から見た左端の条件は、 $S_{r+1} - B_{r} \le S_{l}$ であること。

と書き下すことができます。このように対応付けることで「右端の条件」に含まれる $r$ 関連の要素が $S_{r+1}$ だけに、「左端の条件」に含まれる $l$ 関連の要素が $S_{l}$ だけになるので、累積和の単調性からそれぞれの条件は区間になることが分かります。

このことを利用すると先ほどの例題と（不等号の向き以外は）同じ形に帰着され、解くことができます。

（CSAcademy）FII Code 2020 Round #3 Double Palindromes

CS Academy

この問題では与えられた文字列の連続部分文字列の取り方であって、それが（ある文字列A）＋（Aの反転）＋（A）という形になっているものを数えます。

先ほどと同じ枠組みに帰着させるには、Manacherのアルゴリズムで求められる「回文半径」というものを利用します。そして以下のように対応付けます。

f:id:betrue12:20200328030254p:plain

つまり両端ではなく、途中の区切り位置を $(l, r)$ だと思うことにします。そして両方の回文半径がお互いに「届く」ならば条件を満たす、ということになります。こうすると先ほどの例題に帰着できます。

（JOI）第１４回日本情報オリンピック本選　城壁

E - 城壁 (Rampart)

障害物が存在するグリッド上で、幅1の「正方形の枠」で一定以上の大きさのものを、障害物に邪魔されずに作れる場所の個数を数えます。

この問題は以下のように、斜め45度の線に沿ったマスを1つの列として扱うと見通しが良くなります。

f:id:betrue12:20200328031120p:plain

左端（左上）では、「下と右に障害物に邪魔されずに何マス進めるか」
右端（右下）では、「上と左に障害物に邪魔されずに何マス進めるか」

という値を考えると、上図の矢印が左下・右上ともに出会うことができるなら、障害物に邪魔されずに正方形の枠を作れるということになります。各マスから4方向それぞれに何マス進めるかは $O(HW)$ で全部計算できるので、各斜め線についての処理を例題の形に帰着することができました。実際にはさらに、大きさに関する条件も考慮して答えを計算する必要があります。