WaveletMatrix
(src/data_structure/wavelet_matrix.hpp)

• View this file on GitHub
• Last update: 2024-11-09 01:50:04+09:00
• Include: #include "src/data_structure/wavelet_matrix.hpp"

WaveletMatrix

静的な非負整数列に対するさまざまな区間クエリを $O(\log n)$ で行えるデータ構造です．

コンストラクタ

(1) WaveletMatrix<T> wm(int n)
(2) WaveletMatrix<T> wm(vector<T> a)

• (1): 長さ n の数列 wm を作ります．初期値は $0$ です．
• (2): 長さ n = ssize(v) の数列 wm を作り， v の内容を初期値として build を行います．

計算量

• (1) $O(n)$
• (2) $O(n \log n)$

set

void wm.set(int i, T x)

wm[i] に $x$ を代入します．

制約

• $0 \leq n < n$
• $0 \leq x$
• build をした後に呼び出してはいけない

計算量

• $O(1)$

build

void wm.build()

wm の初期化を行います．

計算量

• $O(n \log n)$

get

T wm.get(int i)

wm[i] を返します．

制約

• $0 \leq i < n$

計算量

• $O(\log n)$

kth_smallest

T wm.kth_smallest(int l, int r, int k)

wm[l, r) における $k$ 番目 ( $0 \mathrm{-indexed}$ ) に小さい値を返します．

制約

• $0 \leq l \leq r \leq n$
• $0 \leq k < r - l$

計算量

• $O(\log n)$

kth_largest

T wm.kth_largest(int l, int r, int k)

wm[l, r) における $k$ 番目 ( $0 \mathrm{-indexed}$ ) に大きい値を返します．

制約

• $0 \leq l \leq r \leq n$
• $0 \leq k < r - l$

計算量

• $O(\log n)$

range_freq

(1) int range_freq(int l, int r, T upper)
(2) int range_freq(int l, int r, T lower, T upper)

(1) wm[l, r) における $\mathrm{upper}$ 未満の要素の個数を返します．
(2) wm[l, r) における $\mathrm{lower}$ 以上 $\mathrm{upper}$ 未満の要素の個数を返します．

制約

• $0 \leq l \leq r \leq n$
• (2) $\mathrm{lower} \leq \mathrm{upper}$

計算量

• $O(\log n)$

prev_value

T wm.prev_value(int l, int r, T upper)

wm[l, r) における $\mathrm{upper}$ 未満の最大値を返します．
そのような値が存在しない存在しない場合は $-1$ を返します．

制約

• $0 \leq l \leq r \leq n$

計算量

• $O(\log n)$

next_value

T wm.next_value(int l, int r, T lower)

wm[l, r) における $\mathrm{lower}$ 以上の最小値を返します．
存在しない場合は $-1$ を返します．

制約

• $0 \leq l \leq r \leq n$

計算量

• $O(\log n)$

Depends on

• template (src/template/template.hpp)

Verified with

• verify/library_checker/data_structure/range_kth_smallest.test.cpp
• verify/library_checker/data_structure/static_range_frequency.test.cpp

Code

#pragma once
#include "../template/template.hpp"
#include <immintrin.h>
struct BitVector {
    using u32 = uint32_t;
    using i64 = int64_t;
    using u64 = uint64_t;
    static constexpr u32 w = 64;
    vector<u64> block;
    vector<u32> count;
    u32 n, zeros;
    inline u32 get(const u32 i) const {
        return u32(block[i / w] >> (i % w)) & 1u;
    }
    inline void set(const u32 i) {
        block[i / w] |= 1LL << (i % w);
    }
    BitVector() {}
    BitVector(const int _n) {
        init(_n);
    }
    __attribute__((optimize("O3,unroll-loops"))) void init(const int _n) {
        n = zeros = _n;
        block.resize(n / w + 1, 0);
        count.resize(block.size(), 0);
    }
    __attribute__((target("popcnt"))) void build() {
        for(u32 i = 1; i < block.size(); ++i) {
            count[i] = count[i - 1] + _mm_popcnt_u64(block[i - 1]);
        }
        zeros = rank0(n);
    }
    inline u32 rank0(const u32 i) const {
        return i - rank1(i);
    }
    __attribute__((target("bmi2,popcnt"))) inline u32 rank1(const u32 i) const {
        return count[i / w] + _mm_popcnt_u64(_bzhi_u64(block[i / w], i % w));
    }
};
template <typename T>
struct WaveletMatrix {
   private:
    using u32 = uint32_t;
    using i64 = int64_t;
    using u64 = uint64_t;
    int n, lg;
    vector<T> a;
    vector<BitVector> bv;
    inline pair<u32, u32> succ0(const int l, const int r, const int h) const {
        return make_pair(bv[h].rank0(l), bv[h].rank0(r));
    }
    inline pair<u32, u32> succ1(const int l, const int r, const int h) const {
        const u32 l0 = bv[h].rank0(l);
        const u32 r0 = bv[h].rank0(r);
        const u32 zeros = bv[h].zeros;
        return make_pair(l + zeros - l0, r + zeros - r0);
    }

   public:
    WaveletMatrix(const u32 _n)
        : n(max<u32>(_n, 1)), a(n) {}
    WaveletMatrix(const vector<T>& _a)
        : n(_a.size()), a(_a) {
        if(n == 0) {
            a.push_back(0);
            n = 1;
        }
        build();
    }
    __attribute__((optimize("O3"))) void build() {
        lg = __lg(max<T>(*max_element(begin(a), end(a)), 1)) + 1;
        bv.assign(lg, n);
        vector<T> cur = a, nxt(n);
        for(int h = lg - 1; h >= 0; --h) {
            for(int i = 0; i < n; ++i)
                if((cur[i] >> h) & 1) bv[h].set(i);
            bv[h].build();
            array<decltype(begin(nxt)), 2> it{begin(nxt), begin(nxt) + bv[h].zeros};
            for(int i = 0; i < n; ++i) *it[bv[h].get(i)]++ = cur[i];
            swap(cur, nxt);
        }
        return;
    }
    void set(const int i, const T& x) {
        assert(0 <= i and i < n);
        assert(x >= 0);
        a[i] = x;
    }
    T get(u32 k) const {
        assert(int(k) < n);
        T ret = 0;
        for(int h = lg - 1; h >= 0; --h) {
            const u32 f = bv[h].get(k);
            ret |= f ? T(1) << h : 0;
            k = f ? bv[h].rank1(k) + bv[h].zeros : bv[h].rank0(k);
        }
        return ret;
    }
    T kth_smallest(u32 l, u32 r, u32 k) const {
        assert(l <= r and int(r) <= n);
        assert(k < r - l);
        T res = 0;
        for(int h = lg - 1; h >= 0; --h) {
            const u32 l0 = bv[h].rank0(l), r0 = bv[h].rank0(r);
            if(k < r0 - l0)
                l = l0, r = r0;
            else {
                k -= r0 - l0;
                res |= (T)1 << h;
                l += bv[h].zeros - l0;
                r += bv[h].zeros - r0;
            }
        }
        return res;
    }
    T kth_largest(const int l, const int r, const int k) const {
        assert(0 <= l and l <= r and r <= n);
        assert(0 <= k and k < r - l);
        return kth_smallest(l, r, r - l - k - 1);
    }
    int range_freq(int l, int r, const T& upper) const {
        assert(0 <= l and l <= r and r <= n);
        if(upper >= (T(1) << lg)) return r - l;
        int ret = 0;
        for(int h = lg - 1; h >= 0; --h) {
            const bool f = (upper >> h) & 1;
            const u32 l0 = bv[h].rank0(l), r0 = bv[h].rank0(r);
            if(f) {
                ret += r0 - l0;
                l += bv[h].zeros - l0;
                r += bv[h].zeros - r0;
            } else {
                l = l0;
                r = r0;
            }
        }
        return ret;
    }
    int range_freq(const int l, const int r, const T& lower, const T& upper) const {
        assert(0 <= l and l <= r and r <= n);
        assert(lower <= upper);
        return range_freq(l, r, upper) - range_freq(l, r, lower);
    }
    T prev_value(const int l, const int r, const T& upper) const {
        assert(0 <= l and l <= r and r <= n);
        const int cnt = range_freq(l, r, upper);
        return cnt == 0 ? T(-1) : kth_smallest(l, r, cnt - 1);
    }
    T next_value(const int l, const int r, const T& lower) const {
        assert(0 <= l and l <= r and r <= n);
        const int cnt = range_freq(l, r, lower);
        return cnt == r - l ? T(-1) : kth_smallest(l, r, cnt);
    }
};

#line 2 "src/template/template.hpp"
#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using P = pair<long long, long long>;
#define rep(i, a, b) for(long long i = (a); i < (b); ++i)
#define rrep(i, a, b) for(long long i = (a); i >= (b); --i)
constexpr long long inf = 4e18;
struct SetupIO {
    SetupIO() {
        ios::sync_with_stdio(0);
        cin.tie(0);
        cout << fixed << setprecision(30);
    }
} setup_io;
#line 3 "src/data_structure/wavelet_matrix.hpp"
#include <immintrin.h>
struct BitVector {
    using u32 = uint32_t;
    using i64 = int64_t;
    using u64 = uint64_t;
    static constexpr u32 w = 64;
    vector<u64> block;
    vector<u32> count;
    u32 n, zeros;
    inline u32 get(const u32 i) const {
        return u32(block[i / w] >> (i % w)) & 1u;
    }
    inline void set(const u32 i) {
        block[i / w] |= 1LL << (i % w);
    }
    BitVector() {}
    BitVector(const int _n) {
        init(_n);
    }
    __attribute__((optimize("O3,unroll-loops"))) void init(const int _n) {
        n = zeros = _n;
        block.resize(n / w + 1, 0);
        count.resize(block.size(), 0);
    }
    __attribute__((target("popcnt"))) void build() {
        for(u32 i = 1; i < block.size(); ++i) {
            count[i] = count[i - 1] + _mm_popcnt_u64(block[i - 1]);
        }
        zeros = rank0(n);
    }
    inline u32 rank0(const u32 i) const {
        return i - rank1(i);
    }
    __attribute__((target("bmi2,popcnt"))) inline u32 rank1(const u32 i) const {
        return count[i / w] + _mm_popcnt_u64(_bzhi_u64(block[i / w], i % w));
    }
};
template <typename T>
struct WaveletMatrix {
   private:
    using u32 = uint32_t;
    using i64 = int64_t;
    using u64 = uint64_t;
    int n, lg;
    vector<T> a;
    vector<BitVector> bv;
    inline pair<u32, u32> succ0(const int l, const int r, const int h) const {
        return make_pair(bv[h].rank0(l), bv[h].rank0(r));
    }
    inline pair<u32, u32> succ1(const int l, const int r, const int h) const {
        const u32 l0 = bv[h].rank0(l);
        const u32 r0 = bv[h].rank0(r);
        const u32 zeros = bv[h].zeros;
        return make_pair(l + zeros - l0, r + zeros - r0);
    }

   public:
    WaveletMatrix(const u32 _n)
        : n(max<u32>(_n, 1)), a(n) {}
    WaveletMatrix(const vector<T>& _a)
        : n(_a.size()), a(_a) {
        if(n == 0) {
            a.push_back(0);
            n = 1;
        }
        build();
    }
    __attribute__((optimize("O3"))) void build() {
        lg = __lg(max<T>(*max_element(begin(a), end(a)), 1)) + 1;
        bv.assign(lg, n);
        vector<T> cur = a, nxt(n);
        for(int h = lg - 1; h >= 0; --h) {
            for(int i = 0; i < n; ++i)
                if((cur[i] >> h) & 1) bv[h].set(i);
            bv[h].build();
            array<decltype(begin(nxt)), 2> it{begin(nxt), begin(nxt) + bv[h].zeros};
            for(int i = 0; i < n; ++i) *it[bv[h].get(i)]++ = cur[i];
            swap(cur, nxt);
        }
        return;
    }
    void set(const int i, const T& x) {
        assert(0 <= i and i < n);
        assert(x >= 0);
        a[i] = x;
    }
    T get(u32 k) const {
        assert(int(k) < n);
        T ret = 0;
        for(int h = lg - 1; h >= 0; --h) {
            const u32 f = bv[h].get(k);
            ret |= f ? T(1) << h : 0;
            k = f ? bv[h].rank1(k) + bv[h].zeros : bv[h].rank0(k);
        }
        return ret;
    }
    T kth_smallest(u32 l, u32 r, u32 k) const {
        assert(l <= r and int(r) <= n);
        assert(k < r - l);
        T res = 0;
        for(int h = lg - 1; h >= 0; --h) {
            const u32 l0 = bv[h].rank0(l), r0 = bv[h].rank0(r);
            if(k < r0 - l0)
                l = l0, r = r0;
            else {
                k -= r0 - l0;
                res |= (T)1 << h;
                l += bv[h].zeros - l0;
                r += bv[h].zeros - r0;
            }
        }
        return res;
    }
    T kth_largest(const int l, const int r, const int k) const {
        assert(0 <= l and l <= r and r <= n);
        assert(0 <= k and k < r - l);
        return kth_smallest(l, r, r - l - k - 1);
    }
    int range_freq(int l, int r, const T& upper) const {
        assert(0 <= l and l <= r and r <= n);
        if(upper >= (T(1) << lg)) return r - l;
        int ret = 0;
        for(int h = lg - 1; h >= 0; --h) {
            const bool f = (upper >> h) & 1;
            const u32 l0 = bv[h].rank0(l), r0 = bv[h].rank0(r);
            if(f) {
                ret += r0 - l0;
                l += bv[h].zeros - l0;
                r += bv[h].zeros - r0;
            } else {
                l = l0;
                r = r0;
            }
        }
        return ret;
    }
    int range_freq(const int l, const int r, const T& lower, const T& upper) const {
        assert(0 <= l and l <= r and r <= n);
        assert(lower <= upper);
        return range_freq(l, r, upper) - range_freq(l, r, lower);
    }
    T prev_value(const int l, const int r, const T& upper) const {
        assert(0 <= l and l <= r and r <= n);
        const int cnt = range_freq(l, r, upper);
        return cnt == 0 ? T(-1) : kth_smallest(l, r, cnt - 1);
    }
    T next_value(const int l, const int r, const T& lower) const {
        assert(0 <= l and l <= r and r <= n);
        const int cnt = range_freq(l, r, lower);
        return cnt == r - l ? T(-1) : kth_smallest(l, r, cnt);
    }
};

Back to top page