Range Kth Smallestに対する4つの解法と2つの実装例

/// Range Kth Smallest (offlineかつstatic)
/// 並列二分探索による解法
/// aの昇順に配列を構築していくと考える。aをどこまで使えば区間[li, ri]にki以上の要素が入るかを考えると、これは単調である。
/// よって、どこまで使うかを二分探索で求めることが出来る。その境界にある要素がまさにRange Kth Smallestの解である。

import std;

void main () {
    int N, Q; readln.read(N, Q);
    auto a = readln.split.to!(int[]);

    auto l = new int[](Q);
    auto r = new int[](Q);
    auto k = new int[](Q);
    foreach (i; 0..Q) readln.read(l[i], r[i], k[i]);

    auto rsq = new SegmentTree!(int, (int a, int b) => a + b, () => 0)(N);
    auto ok = new int[](Q);
    auto ng = new int[](Q);
    auto mid = new int[](Q);
    ok[] = N - 1;
    ng[] = -1;

    auto value_ord = iota(N).array.sort!((i, j) => a[i] < a[j]).array;
    auto query_ord = iota(Q).array;

    while (true) {
        bool end = () {
            foreach (i; 0..Q) if (1 < abs(ok[i] - ng[i])) return false;
            return true;
        }();
        if (end) break;

        // セグメント木のリセット
        foreach (i; 0..N) rsq.set(i, 0);

        // midの計算 + midの昇順にソート
        foreach (i; 0..Q) mid[i] = (ok[i] + ng[i]) / 2;
        query_ord.sort!((i, j) => mid[i] < mid[j]);

        int L = 0, R = 0;
        foreach (count, i; value_ord) {
            // 点追加
            rsq.set(i, rsq.get(i) + 1);

            if (L < Q && count == mid[query_ord[L]]) {
                while (R < Q) {
                    if (mid[query_ord[R]] != count) break;
                    R++;
                }

                // 二分探索を一段階進める
                foreach (U; L..R) {
                    int qidx = query_ord[U];
                    bool f () {
                        return k[qidx] + 1 <= rsq.prod(l[qidx], r[qidx]);
                    }

                    if (f()) {
                        ok[qidx] = mid[qidx];
                    }
                    else {
                        ng[qidx] = mid[qidx];
                    }
                }

                L = R;
            }
        }
    }

    foreach (i; 0..Q) {
        writeln(a[value_ord[ok[i]]]);
    }
}

void read (T...) (string S, ref T args) {
    import std.conv : to;
    import std.array : split;
    auto buf = S.split;
    foreach (i, ref arg; args) {
        arg = buf[i].to!(typeof(arg));
    }
}

// セグメント木は省略

/// Range Kth Smallest (offlineかつstatic)
/// Mo's algorithm + 平方分割rsqによる解法
/// aに含まれる要素で座標圧縮をし、横軸が座標圧縮後のindexとなるrsqを考える。
/// Moによって区間を巡回し、区間内の要素で上のrsqに積んでいくことにすると、[0, x)がはじめてkiを超える場所がRange Kth Smallestの解になる。
/// ここで、rsqに平方分割を採用すると、更新O(1)、ki以上の場所の取得がO(√N)であるから、全体O(Q√N + N√Q)で解ける。

/// 補足1: ki以上の場所がO(√N)になるからくり
/// 先頭から足していくことにして、足して初めてki以上になるブロックを見つける。ここまでO(√N)
/// ブロック内で厳密にkiになる場所を探す。ここでO(√N)
/// 合計O(√N)

/// 補足2: 座圧後のrsqで平方分割を使っているのは、セグメント木を用いるより全体の計算量が良くなるから。Mo's algorithmでは移動が支配的なので、取得を多少遅くしたとしても更新を早くする方が有利になる。

import std;

void main () {
    int N, Q; readln.read(N, Q);
    auto a = readln.split.to!(int[]);

    auto l = new int[](Q);
    auto r = new int[](Q);
    auto k = new int[](Q);
    auto ans = new int[](Q);
    foreach (i; 0..Q) readln.read(l[i], r[i], k[i]);

    // mo's algorithmの準備
    const int mo_width = max(1, sqrt(Q.to!double).to!int);
    auto priority = new int[](Q);
    foreach (i; 0..Q) priority[i] = l[i] / mo_width;

    bool less (int i, int j) {
        if (priority[i] == priority[j]) {
            if (priority[i] % 2 == 0) return r[i] < r[j];
            return r[j] < r[i];
        }
        return priority[i] < priority[j];
    }

    auto query_ord = iota(Q).array.sort!(less).array;

    // 平方分割の準備
    auto f = new int[](N);
    int max_f = 0;
    auto index = iota(N).array.sort!((i, j) => a[i] < a[j]).array;
    auto b = a.dup.sort.uniq.array;
    {
        int L = 0, R = 0;
        while (L < N) {
            while (R < N) {
                if (a[index[L]] != a[index[R]]) break;
                R++;
            }

            foreach (U; L..R) {
                f[index[U]] = max_f;
            }
            max_f++;

            L = R;
        }
    }
    const int width = max(1, sqrt(max_f.to!double).to!int);
    auto block_count = new int[](max_f / width + 1);
    auto count = new int[](max_f + 1);

    int L = 0, R = 0;
    foreach (i; query_ord) {
        // 区間合わせ
        while (L < l[i]) { block_count[f[L] / width]--; count[f[L]]--; L++; }
        while (l[i] < L) { L--; block_count[f[L] / width]++; count[f[L]]++; }
        while (R < r[i]) { block_count[f[R] / width]++; count[f[R]]++; R++; }
        while (r[i] < R) { R--; block_count[f[R] / width]--; count[f[R]]--; }

        // 取得
        int acc = 0;
        int cur_block = 0;
        while (acc + block_count[cur_block] < k[i] + 1) {
            acc += block_count[cur_block];
            cur_block++;
        }

        // ブロック内走査
        int cur = cur_block * width;
        while (acc + count[cur] < k[i] + 1) {
            acc += count[cur];
            cur++;
        }

        ans[i] = b[cur];
    }

    foreach (i; 0..Q) {
        writeln(ans[i]);
    }
}

void read (T...) (string S, ref T args) {
    import std.conv : to;
    import std.array : split;
    auto buf = S.split;
    foreach (i, ref arg; args) {
        arg = buf[i].to!(typeof(arg));
    }
}