Wavelet Matrix
(structure/WaveletMatrix.cpp)

• View this file on GitHub
• Last update: 2021-09-12 23:44:25+09:00

概要

Wavelet Matrix

• rank(c,x) : $[0,c)$ に x が何回存在するか
• quantile(l,r,c) : $[l,r)$ 中で $c+1$ 番目に大きい数

計算量

全て $O(log \sigma)$ ($\sigma$は数の種類数)

Depends on

• Bit Vector (structure/BitVector.cpp)
• template/template.cpp

Required by

• FM Index (string/FM_index.cpp)

Verified with

• test/FM_index.test.cpp
• test/WaveletMatrix.test.cpp

Code

#pragma once
#include "../template/template.cpp"
#include "BitVector.cpp"

template <class T, class C>
class WaveletMatrix {
    int N, bitlen;
    vector<BitVector> index;
    map<C, int> st;

   public:
    T body;
    int rank(C c, int idx) {
        if (st.find(c) == st.end()) return 0;
        rev(i, bitlen) {
            if (c >> i & 1)
                idx = index[i].rank(1, idx) + index[i].rank(0, N);
            else
                idx -= index[i].rank(1, idx);
        }
        return max(0, idx - st[c]);
    }
    C quantile(int l, int r, int c) {
        C res = 0;
        rev(i, bitlen) {
            ll cnt = (r - l) - (index[i].rank(1, r) - index[i].rank(1, l));
            if (cnt <= c) {
                c -= cnt;
                l = index[i].rank(0, N) + index[i].rank(1, l);
                r = index[i].rank(0, N) + index[i].rank(1, r);
                res += 1ll << i;
            } else {
                l -= index[i].rank(1, l);
                r -= index[i].rank(1, r);
            }
        }
        return res;
    }
    WaveletMatrix(T V, ll bitlen) : N(len(V)), bitlen(bitlen), body(V) {
        vector<bool> bit(N);
        index.resize(bitlen, bit);
        rev(i, bitlen) {
            T newV[2];
            rep(j, N) {
                bit[j] = (V[j] >> i & 1);
                newV[V[j] >> i & 1].push_back(V[j]);
            }
            V = newV[0];
            V.insert(V.end(), all(newV[1]));
            index[i] = BitVector(bit);
        }
        rep(i, N) if (st.find(V[i]) == st.end()) st[V[i]] = i;
    }
};
/*
@brief Wavelet Matrix
@docs docs/WaveletMatrix.md
*/

#line 2 "template/template.cpp"
#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define rep(i, n) for (int i = 0; i < n; i++)
#define REP(i, n) for (int i = 1; i < n; i++)
#define rev(i, n) for (int i = n - 1; i >= 0; i--)
#define REV(i, n) for (int i = n - 1; i > 0; i--)
#define all(v) v.begin(), v.end()
#define PL pair<ll, ll>
#define PI pair<int, int>
#define pi acos(-1)
#define len(s) (int)s.size()
#define compress(v) \
    sort(all(v));   \
    v.erase(unique(all(v)), v.end());
#define comid(v, x) lower_bound(all(v), x) - v.begin()

template<class T>
using prique=priority_queue<T,vector<T>,greater<>>;

template <class T, class U>
inline bool chmin(T &a, U b) {
    if (a > b) {
        a = b;
        return true;
    }
    return false;
}
template <class T, class U>
inline bool chmax(T &a, U b) {
    if (a < b) {
        a = b;
        return true;
    }
    return false;
}
constexpr ll inf = 3e18;
#line 3 "structure/BitVector.cpp"

class BitVector {
    vector<ll> sum;
    vector<uint64_t> bit;

   public:
    ll rank(bool val, ll idx) {
        uint64_t mask = ((uint64_t)1 << (idx & ((1 << 6) - 1))) - 1;
        ll res = sum[idx >> 6] + __builtin_popcountll(bit[idx >> 6] & mask);
        return (val ? res : idx - res);
    }
    BitVector(vector<bool>& v) {
        ll sz = (len(v) >> 6) + 1;
        bit.assign(sz, 0);
        sum.assign(sz, 0);
        rep(i, len(v)) {
            bit[i >> 6] |= (uint64_t)(v[i]) << (i & ((1 << 6) - 1));
        }
        rep(i, sz - 1) {
            sum[i + 1] = sum[i] + __builtin_popcountll(bit[i]);
        }
    }
};
/*
@brief Bit Vector
@docs docs/BitVector.md
*/
#line 4 "structure/WaveletMatrix.cpp"

template <class T, class C>
class WaveletMatrix {
    int N, bitlen;
    vector<BitVector> index;
    map<C, int> st;

   public:
    T body;
    int rank(C c, int idx) {
        if (st.find(c) == st.end()) return 0;
        rev(i, bitlen) {
            if (c >> i & 1)
                idx = index[i].rank(1, idx) + index[i].rank(0, N);
            else
                idx -= index[i].rank(1, idx);
        }
        return max(0, idx - st[c]);
    }
    C quantile(int l, int r, int c) {
        C res = 0;
        rev(i, bitlen) {
            ll cnt = (r - l) - (index[i].rank(1, r) - index[i].rank(1, l));
            if (cnt <= c) {
                c -= cnt;
                l = index[i].rank(0, N) + index[i].rank(1, l);
                r = index[i].rank(0, N) + index[i].rank(1, r);
                res += 1ll << i;
            } else {
                l -= index[i].rank(1, l);
                r -= index[i].rank(1, r);
            }
        }
        return res;
    }
    WaveletMatrix(T V, ll bitlen) : N(len(V)), bitlen(bitlen), body(V) {
        vector<bool> bit(N);
        index.resize(bitlen, bit);
        rev(i, bitlen) {
            T newV[2];
            rep(j, N) {
                bit[j] = (V[j] >> i & 1);
                newV[V[j] >> i & 1].push_back(V[j]);
            }
            V = newV[0];
            V.insert(V.end(), all(newV[1]));
            index[i] = BitVector(bit);
        }
        rep(i, N) if (st.find(V[i]) == st.end()) st[V[i]] = i;
    }
};
/*
@brief Wavelet Matrix
@docs docs/WaveletMatrix.md
*/

Back to top page