Day 085-Q1 — Wavelet Matrix（区間 $k$ 番目・rank/select・range_freq $O(\log \sigma)$）

問題

長さ $N$ の非負整数列 $A = (a_0, \dots, a_{N-1})$ が与えられる。$Q$ 個のクエリを処理せよ。

• 1 l r k: 区間 $[l, r)$ を昇順に並べたときの $k$ 番目（0-indexed）の値（quantile / k-th smallest）。
• 2 l r x: 区間 $[l, r)$ に含まれる、値が $x$ 未満の要素の個数（range freq）。

Wavelet Matrix を用いて各クエリを $O(\log \sigma)$（$\sigma$ は値域）で処理せよ。

入出力例

5 3
3 1 4 1 5
1 0 5 2
2 0 5 4
1 1 4 0

3
3
1

$[0,5)$ ソート $[1,1,3,4,5]$ の2番目は 3。値 4 未満は 3 個。$[1,4)=[1,4,1]$ の0番目は 1。

概念図: Wavelet Matrix の層構造と quantile 分岐

ヒント

def quantile(l, r, k):
    res = 0
    for b in reversed(range(BITS)):
        c0 = rank0(b, r) - rank0(b, l)
        if k < c0:
            l, r = rank0(b, l), rank0(b, r)
        else:
            k -= c0; res |= (1 << b)
            l = zeros[b] + rank1(b, l)
            r = zeros[b] + rank1(b, r)
    return res

模範解答

import sys
input = sys.stdin.readline

class BitVector:
    __slots__ = ('n', 'block', 'bit')
    def __init__(self, bits):
        self.n = len(bits)
        nb = (self.n >> 5) + 1
        self.bit = [0] * nb
        self.block = [0] * (nb + 1)
        for i, b in enumerate(bits):
            if b:
                self.bit[i >> 5] |= 1 << (i & 31)
        for i in range(nb):
            self.block[i + 1] = self.block[i] + bin(self.bit[i]).count('1')
    def rank1(self, x):
        return self.block[x >> 5] + bin(self.bit[x >> 5] & ((1 << (x & 31)) - 1)).count('1')
    def rank0(self, x):
        return x - self.rank1(x)

class WaveletMatrix:
    def __init__(self, arr, bits):
        self.n = len(arr)
        self.bits = bits
        self.mat = []
        self.zeros = [0] * bits
        cur = arr[:]
        for b in reversed(range(bits)):
            bv = BitVector([(v >> b) & 1 for v in cur])
            self.mat.append(bv)
            self.zeros[b] = bv.rank0(self.n)
            zero = [v for v in cur if not ((v >> b) & 1)]
            one = [v for v in cur if (v >> b) & 1]
            cur = zero + one

    def _bv(self, b):
        return self.mat[self.bits - 1 - b]

    def quantile(self, l, r, k):
        res = 0
        for b in reversed(range(self.bits)):
            bv = self._bv(b)
            c0 = bv.rank0(r) - bv.rank0(l)
            if k < c0:
                l, r = bv.rank0(l), bv.rank0(r)
            else:
                k -= c0
                res |= (1 << b)
                l = self.zeros[b] + bv.rank1(l)
                r = self.zeros[b] + bv.rank1(r)
        return res

    def range_freq(self, l, r, x):
        if x <= 0: return 0
        if x >= (1 << self.bits): return r - l
        res = 0
        for b in reversed(range(self.bits)):
            bv = self._bv(b)
            if (x >> b) & 1:
                res += bv.rank0(r) - bv.rank0(l)
                l = self.zeros[b] + bv.rank1(l)
                r = self.zeros[b] + bv.rank1(r)
            else:
                l, r = bv.rank0(l), bv.rank0(r)
        return res

def solve():
    N, Q = map(int, input().split())
    A = list(map(int, input().split()))
    wm = WaveletMatrix(A, 18)
    out = []
    for _ in range(Q):
        t, l, r, x = map(int, input().split())
        out.append(wm.quantile(l, r, x) if t == 1 else wm.range_freq(l, r, x))
    print('\n'.join(map(str, out)))

solve()

計算量: 構築 $O(N \log \sigma)$、各クエリ $O(\log \sigma)$。

Step-by-Step 解説

Step 1: succinct ビットベクトルの rank

32bit ブロックごとに popcount の累積和を持ち、rank1(x) はブロック累積 + 端数ワードのマスク popcount で $O(1)$。

Step 2: 安定ソートによる層構築

最上位ビットから「0 の要素を前・1 の要素を後」に安定分割していく。この写像が rank だけで表現できる。

Step 3: quantile の分岐

Step 4: range_freq

$x$ のビットが 1 の層では「0 側は全て $x$ 未満確定」で加算し 1 側へ潜る。ビット 0 の層では 0 側へ潜るだけ。

よくあるミス

次のステップ

• 発展問題: 区間 top-k 総和（Wavelet Matrix + 重み累積和）
• 発展問題: 2次元点の矩形内 k 番目（座標圧縮 + Wavelet Matrix）

自己評価