Day 127-Q1 | Wavelet Matrix

問題

長さ $N$ の数列 $A_1,\dots,A_N$（$0 \le A_i < 2^{18}$）が与えられる。$Q$ 個のクエリに答えよ。各クエリは $L,R,K$ で与えられ、区間 $[L,R]$（$1$-indexed、両端含む）に含まれる要素を昇順に並べたとき、$K$番目に小さい値を出力せよ。

入力形式

N Q
A_1 A_2 ... A_N
L_1 R_1 K_1
...
L_Q R_Q K_Q

制約

$1 \le N, Q \le 200000$

$0 \le A_i < 2^{18}$

$1 \le L_i \le R_i \le N$

$1 \le K_i \le R_i-L_i+1$

入出力例

入力例1

6 2
5 3 8 1 9 2
2 5 3
1 6 1

出力例1

8
1

A_2..A_5=(3,8,1,9)を昇順にすると(1,3,8,9)、3番目は8。全体の1番目は1。

概念図

ヒント（段階的開示）

ヒント1（方向性）

クエリごとに区間を切り出してソートすると $O(QN\log N)$ になり間に合わない。今回は配列全体に対して一度だけ前処理を行い、どんな区間にも共通のテーブルだけで答えられる構造を作る。これが Wavelet Matrix である。

ヒント2（アプローチ）

値を上位ビットから下位ビットへ見ていき、各レベルで「ビット0の要素を前、1の要素を後ろ」に安定に分割する。各レベルで区間内のビット0の個数を答えられる rank0 前計算配列を持てば、k番目小さい値のクエリを $O(\log(\max A))$ で解ける。

ヒント3（誘導）

zeros_in_range = rank0[level][r] - rank0[level][l]
if k <= zeros_in_range:
    l, r = rank0[level][l], rank0[level][r]
    bit = 0
else:
    k -= zeros_in_range
    zc = zero_count[level]
    l, r = zc + (l - rank0[level][l]), zc + (r - rank0[level][r])
    bit = 1
ans = (ans << 1) | bit

模範解答 (Python)

import sys


def solve():
    data = sys.stdin.buffer.read().split()
    idx = 0
    n = int(data[idx]); idx += 1
    q = int(data[idx]); idx += 1
    a = [int(data[idx + i]) for i in range(n)]
    idx += n

    BIT = 18

    order = a[:]
    rank0 = [None] * BIT
    zero_count = [0] * BIT

    for level in range(BIT - 1, -1, -1):
        bits = [(v >> level) & 1 for v in order]
        pre = [0] * (n + 1)
        for i in range(n):
            pre[i + 1] = pre[i] + (1 - bits[i])
        rank0[level] = pre
        zero_count[level] = pre[n]

        zeros = []
        ones = []
        for v, b in zip(order, bits):
            (zeros if b == 0 else ones).append(v)
        order = zeros + ones

    out = []
    for _ in range(q):
        L = int(data[idx]); idx += 1
        R = int(data[idx]); idx += 1
        K = int(data[idx]); idx += 1

        l, r = L - 1, R
        k = K
        ans = 0
        for level in range(BIT - 1, -1, -1):
            pre = rank0[level]
            zeros_in_range = pre[r] - pre[l]
            if k <= zeros_in_range:
                l, r = pre[l], pre[r]
                ans <<= 1
            else:
                k -= zeros_in_range
                zc = zero_count[level]
                l, r = zc + (l - pre[l]), zc + (r - pre[r])
                ans = (ans << 1) | 1
        out.append(str(ans))

    print('\n'.join(out))


solve()

Step-by-Step 解説

1ビット平面ごとの安定分割
最上位ビットから見て「ビット0」「ビット1」に分け、元の相対順序を保ったまま連結し直す。全ビット分繰り返すと値の昇順に近い構造になる。

2各レベルのrank0配列
各レベルで「先頭からi個中のビット0の個数」を表す前計算配列を持てば、任意区間の0の個数が $O(1)$ で求まる。

3クエリ処理
最上位ビットから「区間内のビット0の個数」を調べ、K番目がそちらにあるか判定して範囲を絞り込む。

4区間の変換規則
0側は pre[l], pre[r] がそのまま新しい範囲。1側は zero_count[level] をオフセットし、1側内での相対位置を使う。

5計算量
構築 $O(N\log(\max A))$、クエリ $O(\log(\max A))$、全体 $O((N+Q)\log(\max A))$。

よくあるミス

ミス	原因	正しい書き方
1側への変換で `r` の更新に `pre[l]` を使う	lとrを混同	`zc+(l-pre[l]), zc+(r-pre[r])`のように対応させる
ビット幅BITを値の上限に対して不足させる	A_iの上限を過小評価	$A_i<2^{18}$ なら BIT=18 を確保する
区間を閉区間のまま扱う	[L,R]と半開[l,r)の変換忘れ	最初に `l,r=L-1,R` で半開区間に統一
安定分割を無視して単純ソートする	元の位置対応が崩れる	走査しながら0リスト・1リストに振り分けて連結

次のステップ

発展: rank0をpopcountブロック単位で持つ真のsuccinct実装にすると定数倍が大幅改善する
次回予告: MSTの辺分類（Kruskal法 + サイクル性質によるmust-use/may-use/never-use判定）

自己評価

完全理解おおむね理解要復習

Day 127-Q1 — Wavelet Matrix（ウェーブレット行列）

問題