Wavelet Matrix
(data-structure/wavelet_matrix.hpp)

• View this file on GitHub
• Last update: 2024-01-08 02:22:28+09:00
• Include: #include "data-structure/wavelet_matrix.hpp"

Description

Wavelet 行列は，静的な非負整数列に対する様々なクエリが処理できるデータ構造である．非負整数列を上位ビットからソートして完備辞書を構築する．

空間計算量: $O(n \log m)$, $m$ は数列およびクエリで与えられる非負整数の最大値

Operations

• WaveletMatrix(vector<T> v)
  • v からウェーブレット行列を構築する
  • 時間計算量: $O(n \log m)$
• T operator[](int k)
  • $k$ 番目の要素を取得する
  • 時間計算量: $O(\log m)$
• int rank(int k, T x)
  • $[0, k)$ での $x$ の個数を返す
  • 時間計算量: $O(\log m)$
• int range_freq(int l, int r, T ub)
  • $[l, r)$ で値が $[0, ub)$ に含まれる要素の個数を返す
  • 時間計算量: $O(\log m)$
• int range_freq(int l, int r, T lb, T ub)
  • $[l, r)$ で値が $[lb, ub)$ に含まれる要素の個数を返す
  • 時間計算量: $O(\log m)$
• int select(int k, T x)
  • $k$ 番目の $x$ の位置を返す
  • 時間計算量: $O(\log n \log m)$
• T quantile(int l, int r, int k)
  • $[l, r)$ で $k$ 番目に小さい値を返す
  • 時間計算量: $O(\log m)$

Reference

• ウェーブレット行列(wavelet matrix)

Depends on

• Rank/Select Dictionary (data-structure/bit_vector.hpp)

Verified with

• test/yosupo/range_kth_smallest.test.cpp
• test/yosupo/static_range_frequency.test.cpp

Code

#pragma once
#include <algorithm>
#include <unordered_map>
#include <vector>

#include "bit_vector.hpp"

template <typename T, int MAX_LOG>
class WaveletMatrix {
   public:
    WaveletMatrix() = default;
    explicit WaveletMatrix(std::vector<T> v) : mat(MAX_LOG), cnt0(MAX_LOG) {
        const int n = v.size();
        for (int d = MAX_LOG - 1; d >= 0; --d) {
            std::vector<bool> bit(n);
            for (int i = 0; i < n; ++i) bit[i] = v[i] >> d & 1;
            mat[d] = BitVector(bit);
            std::vector<int> cnt(2);
            for (int i = 0; i < n; ++i) {
                if (!bit[i]) cnt[0]++;
            }
            cnt0[d] = cnt[0];
            cnt[1] = n;
            std::vector<T> nv(n);
            for (int i = n - 1; i >= 0; --i) {
                nv[--cnt[bit[i]]] = v[i];
            }
            v.swap(nv);
        }
        for (int i = 0; i < n; ++i) {
            if (!start.contains(v[i])) start[v[i]] = i;
        }
    }

    T operator[](int k) const {
        T ret = 0;
        for (int d = mat.size() - 1; d >= 0; --d) {
            bool b = mat[d][k];
            ret |= T(b) << d;
            k = cnt0[d] * b + mat[d].rank(k, b);
        }
        return ret;
    }

    int rank(int k, T x) const {
        for (int d = mat.size() - 1; d >= 0; --d) {
            bool b = x >> d & 1;
            k = cnt0[d] * b + mat[d].rank(k, b);
        }
        if (start.contains(x)) return k - start.at(x);
        return k;
    }

    int range_freq(int l, int r, T ub) const {
        if (l >= r) return 0;
        int ret = 0;
        for (int d = mat.size() - 1; d >= 0; --d) {
            bool b = ub >> d & 1;
            if (b) ret += mat[d].rank(r, 0) - mat[d].rank(l, 0);
            l = cnt0[d] * b + mat[d].rank(l, b);
            r = cnt0[d] * b + mat[d].rank(r, b);
        }
        return ret;
    }

    int range_freq(int l, int r, T lb, T ub) const {
        if (lb >= ub) return 0;
        return range_freq(l, r, ub) - range_freq(l, r, lb);
    }

    int select(int k, T x) const {
        k += start.at(x);
        for (int d = 0; d < (int)mat.size(); ++d) {
            bool b = x >> d & 1;
            k = mat[d].select(k - cnt0[d] * b, b);
        }
        return k;
    }

    T quantile(int l, int r, int k) const {
        T ret = 0;
        for (int d = (int)mat.size() - 1; d >= 0; --d) {
            int cnt = mat[d].rank(r, 0) - mat[d].rank(l, 0);
            bool b = k < cnt ? 0 : 1;
            l = cnt0[d] * b + mat[d].rank(l, b);
            r = cnt0[d] * b + mat[d].rank(r, b);
            if (b) {
                ret |= T(1) << d;
                k -= cnt;
            }
        }
        return ret;
    }

   private:
    std::vector<BitVector> mat;
    std::vector<int> cnt0;
    std::unordered_map<T, int> start;
};

#line 2 "data-structure/wavelet_matrix.hpp"
#include <algorithm>
#include <unordered_map>
#include <vector>

#line 2 "data-structure/bit_vector.hpp"
#include <bit>
#include <cstdint>
#include <numeric>
#line 6 "data-structure/bit_vector.hpp"

class BitVector {
   public:
    BitVector() = default;
    explicit BitVector(const std::vector<bool>& v) {
        const int n = v.size() / sz + 1;
        data.resize(n);
        sum.resize(n + 1);
        for (int i = 0; i < (int)v.size(); ++i)
            data[i / sz] |= v[i] << (i % sz);
        for (int i = 0; i < n; ++i)
            sum[i + 1] = sum[i] + std::popcount(data[i]);
    }

    bool operator[](int k) const { return data[k / sz] >> (k % sz) & 1; }

    int rank(int k, bool b) const {
        int mask = (1U << (k % sz)) - 1;
        int r = sum[k / sz] + std::popcount(data[k / sz] & mask);
        return b ? r : k - r;
    }

    int select(int k, bool b) const {
        int lb = 0, ub = data.size();
        while (ub - lb > 1) {
            int m = std::midpoint(lb, ub);
            if (rank(m, b) <= k)
                lb = m;
            else
                ub = m;
        }
        return lb;
    }

   private:
    static constexpr int sz = 32;

    std::vector<uint32_t> data;
    std::vector<int> sum;
};
#line 7 "data-structure/wavelet_matrix.hpp"

template <typename T, int MAX_LOG>
class WaveletMatrix {
   public:
    WaveletMatrix() = default;
    explicit WaveletMatrix(std::vector<T> v) : mat(MAX_LOG), cnt0(MAX_LOG) {
        const int n = v.size();
        for (int d = MAX_LOG - 1; d >= 0; --d) {
            std::vector<bool> bit(n);
            for (int i = 0; i < n; ++i) bit[i] = v[i] >> d & 1;
            mat[d] = BitVector(bit);
            std::vector<int> cnt(2);
            for (int i = 0; i < n; ++i) {
                if (!bit[i]) cnt[0]++;
            }
            cnt0[d] = cnt[0];
            cnt[1] = n;
            std::vector<T> nv(n);
            for (int i = n - 1; i >= 0; --i) {
                nv[--cnt[bit[i]]] = v[i];
            }
            v.swap(nv);
        }
        for (int i = 0; i < n; ++i) {
            if (!start.contains(v[i])) start[v[i]] = i;
        }
    }

    T operator[](int k) const {
        T ret = 0;
        for (int d = mat.size() - 1; d >= 0; --d) {
            bool b = mat[d][k];
            ret |= T(b) << d;
            k = cnt0[d] * b + mat[d].rank(k, b);
        }
        return ret;
    }

    int rank(int k, T x) const {
        for (int d = mat.size() - 1; d >= 0; --d) {
            bool b = x >> d & 1;
            k = cnt0[d] * b + mat[d].rank(k, b);
        }
        if (start.contains(x)) return k - start.at(x);
        return k;
    }

    int range_freq(int l, int r, T ub) const {
        if (l >= r) return 0;
        int ret = 0;
        for (int d = mat.size() - 1; d >= 0; --d) {
            bool b = ub >> d & 1;
            if (b) ret += mat[d].rank(r, 0) - mat[d].rank(l, 0);
            l = cnt0[d] * b + mat[d].rank(l, b);
            r = cnt0[d] * b + mat[d].rank(r, b);
        }
        return ret;
    }

    int range_freq(int l, int r, T lb, T ub) const {
        if (lb >= ub) return 0;
        return range_freq(l, r, ub) - range_freq(l, r, lb);
    }

    int select(int k, T x) const {
        k += start.at(x);
        for (int d = 0; d < (int)mat.size(); ++d) {
            bool b = x >> d & 1;
            k = mat[d].select(k - cnt0[d] * b, b);
        }
        return k;
    }

    T quantile(int l, int r, int k) const {
        T ret = 0;
        for (int d = (int)mat.size() - 1; d >= 0; --d) {
            int cnt = mat[d].rank(r, 0) - mat[d].rank(l, 0);
            bool b = k < cnt ? 0 : 1;
            l = cnt0[d] * b + mat[d].rank(l, b);
            r = cnt0[d] * b + mat[d].rank(r, b);
            if (b) {
                ret |= T(1) << d;
                k -= cnt;
            }
        }
        return ret;
    }

   private:
    std::vector<BitVector> mat;
    std::vector<int> cnt0;
    std::unordered_map<T, int> start;
};

Back to top page