憧れの「セグ木上の二分探索」の使い方をABC382-Cで覚えよう

以下、セグ木とは、ACL（AtCoder Library）のセグ木を使うことを想定としています

ACL の SegTree のドキュメントを見て普通にわかる人には特に参考にはなりません

　セグ木（セグメントツリー）を一応使える人は増えましたが、†セグ木上の二分探索† はまだ使ったことがないという人も多いのではないでしょうか。これを使うといい感じに問題が出たので、使い方を含めて解説したいと思います。

累積 min を二分探索で解く手法

　前提としてわかりやすいので、まず、累積 min を定義してから二分探索で解く手法を説明します。（計算量は全く同じです、たぶん）

　問題のサンプル $3$ で考えます。

　問題の条件から、「自分より美食度が低い人」が自分より前にいた場合、その人は何も食べられません。なので、有効な美食度は累積 min となり、広義単調減少になります。

10 5
60 83 76 45 70 91 37 58 94 22
70 39 52 33 18

　サンプル $3$ がこうなので、美食度の累積 min は以下のようになります。

2147483647 60 60 60 45 45 45 37 37 37 22

　$N$ 個の配列には、$N+1$ 個の「仕切り」が必要なので、累積 min の配列の要素数は $N+1$ であることに注意してください。

　先頭の 2147483647 は min の 単位元 となります。何と作用させても演算結果が相手の値になるようなものです。最小値に対しては、無限に大きい値がそのような性質を持ちますが、無限の値は持てないので、実用上問題ない程度に大きい値（例えば $2^{31}-1$ や $10^9$ ）がその代用として使われます。

　ある寿司は、その美味しさ 以下の 累積 min を持つ中で 最も左 の人に食べられるので、後は二分探索をすればよいです。

セグ木上の二分探索

　実は、セグ木の二分探索でやることはこれとほぼ同じになります。ここからは、コードと一緒に説明した方がわかりやすいので、そうします。

#include <atcoder/segtree>
#include <iostream>
#include <vector>
#include <climits>

using namespace std;
using namespace atcoder;

// 演算 op: 最小値を返す関数
int op(int a, int b) {
    return min(a, b);
}

// 単位元 e: 最小値を求めるときの初期値
int e() {
    return INT_MAX;
}

// 条件 f: prod > target を満たすか
int target;
bool f(int prod) {
    return prod > target;
}

int main() {
    int n, m;
    cin >> n >> m;
    
    vector<int> a(n), b(m);
    for (int i = 0; i < n; ++i) {
        cin >> a[i];
    }
    for (int i = 0; i < m; ++i) {
        cin >> b[i];
    }
    
    // セグメントツリーの初期化
    segtree<int, op, e> seg(a);
    
    // 答えを格納するためのベクター
    vector<int> ans(m);
    for (int i = 0; i < m; ++i) {
        target = b[i];
        // 条件を満たす最大の右端のインデックスを求める
        int idx = seg.max_right<f>(0);
        if (idx < n) {
            cout << idx + 1 << endl; // 1-indexed の出力
        } else {
            cout << -1 << endl; // 条件を満たすものがない場合
        }
    }

    return 0;
}

　まず、配列をセグ木を格納した時点で、前節の累積 min 配列と全く同じ情報が得られます。

    for(int i=0;i<n+1;i++){
        cout << seg.prod(0,i) << endl;
    }

　仮にこういうコードを挟んだ場合、

2147483647
60
60
60
45
45
45
37
37
37
22

　こういう結果が得られて、前節の累積 min とまったく同じになります。

　さて、この時点で、seg.prod を使えば二分探索で答えを絞り込めることがわかります。しかし、seg.prod 自体に $\text{log}$ がかかるので、それを用いて二分探索をすると、 $\text{log}$ が $2$ つ付いてしまうことになります。それでも間に合う場合は多いんですが。

　ここで活躍するのが セグ木上の二分探索 です。これだと $\text{log}$ が $1$ つで済みます。ACL だと seg.max_right および seg.min_left というのがそれにあたります。

　具体的なコードでは、

        int idx = seg.max_right<f>(0);

　という部分がそれにあたります。「$0$ の位置の仕切りからはじめて、条件を満たすだけ右に行ったとこにある仕切りの位置」がこれの求めているものになります。

　ここで、一番わかりづらいのは <f> の部分だと思います。これは、比較関数 を定義しています。これが TRUE か FALSE かによって二分探索を行います。そして、この f は、

// 演算 op: 最小値を返す関数
int op(int a, int b) {
    return min(a, b);
}

// 単位元 e: 最小値を求めるときの初期値
int e() {
    return INT_MAX;
}

// 条件 f: prod > target を満たすか
int target;
bool f(int prod) {
    return prod > target;
}

　ここで定義をしています。ACL のドキュメントだと v となっているところは、prod とした方がわかりやすいと思うので、そうしています。上のコードでは、seg.prod が target より上である時に TRUE を返します。

　これを踏まえて、seg.max_right<f>(0) のお気持ちをより詳しく述べると、「$0$ の位置の仕切りからはじめて、seg.prod(0,R) の値が target より上であるような最も右の R 」となります。

　これは、累積 min において、target 以下であるような最も左の仕切りと一致します。

　やっていることの本質や、計算量は、累積 min を取る解法と変わりませんが、セグ木の扱いになれている場合は、こちらの方が書きやすいという人はいるかもしれません。また、累積となる配列を別に定義しなくてよいので、$N$ とか $N+1$ がごっちゃにならないのもセグ木を使うメリットの $1$ つかなと、個人的には思います。

おまけ（別解：貪欲）

　公式解法とも二分探索とも違う解法を紹介します。

　寿司は、どのような順番で流れてこようと全く同じ結果になります。そこで、食べる人の気持ちになって考えると、美味しい寿司の方が先に来てくれる方が探しやすいです（先頭から食べたいだけ食べて、自分の舌を満足させられない寿司はまとめて後ろに流せるのですから）。よって、寿司を美味しい順にソートすれば、余計な行き来をすることなく探索できます。

　本番中の提出なので整形していないですが、参考までに。