Wavelet Matrix
(data_structure/WaveletMatrix.hpp)

• View this file on GitHub
• Last update: 2023-12-28 15:52:36+09:00
• Include: #include "data_structure/WaveletMatrix.hpp"
• Link: https://miti-7.hatenablog.com/entry/2018/04/28/152259

説明

整数列$(a_0, a_1, \dots, a_{n-1})$を扱う静的なデータ構造である.

コンストラクタ

WaveletMatrix<T> wm(std::vector<T> v);

• 整数型$T$
• 制数列$(a_0, a_1, \dots, a_{n-1})$

を用いる.

Operator

• wm.access(int i)
  • $a_i$を返す
  • 時間計算量$O(\log \sigma)$
• wm.rank(int r, T x)
  • $[0,r)$に含まれる$x$の数を返す.
  • 時間計算量$O(\log \sigma)$
• wm.select(T x, int k)
  • 整数列$(a_0, a_1, \dots, a_{n-1})$に対して$x$が$k$番目に現れる$index$を返す
  • 時間計算量$O(\log \sigma)$
• wm.quantile(int l, int r, int k)
  • $[l,r)$に含まれる要素について$k$番目に小さい値を返す
  • 時間計算量$O(\log \sigma)$

Depends on

• data_structure/bitVector.hpp
• template/int_alias.hpp

Verified with

• test/data_structure/Range_Kth_Smallest.test.cpp
• test/data_structure/Static_Range_Frequency.test.cpp

Code

#pragma once

#include "../data_structure/bitVector.hpp"


/*
    reference: https://miti-7.hatenablog.com/entry/2018/04/28/152259
*/

#include <limits>

#include <map>

#include <vector>


namespace ebi {

template <class T> struct WaveletMatrix {
  private:
    int wordsize = std::numeric_limits<T>::digits;
    int n;
    std::vector<bitVector> mat;
    std::vector<int> border;
    std::map<T, int> map;

  public:
    WaveletMatrix(std::vector<T> v)
        : n(v.size()), mat(wordsize, bitVector(n)), border(wordsize) {
        for (int bit = wordsize - 1; bit >= 0; --bit) {
            std::vector<T> one, zero;
            zero.reserve(n);
            for (int i = 0; i < n; ++i) {
                if (v[i] & ((T)1 << bit)) {
                    one.emplace_back(v[i]);
                    mat[bit].set(i, 1);
                } else {
                    zero.emplace_back(v[i]);
                }
            }
            border[bit] = zero.size();
            mat[bit].build();
            std::copy(one.begin(), one.end(), std::back_inserter(zero));
            v = zero;
        }
        for (int i = 0; i < n; i++) {
            if (map[v[i]] == 0) {
                map[v[i]] = i;
            }
        }
    }

    T access(int i) {
        T val = 0;
        int k = i;
        for (int bit = wordsize - 1; bit >= 0; --bit) {
            if (mat[bit].access(k) == 1) {
                val |= ((T)1 << bit);
                k = border[bit] + mat[bit].rank(k);
            } else {
                k = k - mat[bit].rank(k);
            }
        }
        return val;
    }

    int rank(int r, T x) {
        int k = r;
        for (int bit = wordsize - 1; bit >= 0; --bit) {
            if (x & ((T)1 << bit)) {
                k = border[bit] + mat[bit].rank(k);
            } else {
                k = k - mat[bit].rank(k);
            }
        }
        return k - map[x];
    }

    int select(T x, int k) {
        k = map[x] + k;
        for (int bit = 0; bit < wordsize; ++bit) {
            if (x & ((T)1 << bit)) {
                k = mat[bit].select(k - border[bit]);
            } else {
                k = mat[bit].select0(k);
            }
        }
        return k - 1;
    }

    T quantile(int l, int r, int k) {
        T val = 0;
        for (int bit = wordsize - 1; bit >= 0; --bit) {
            int rank_l = mat[bit].rank(l);
            int rank_r = mat[bit].rank(r);
            int one = rank_r - rank_l;
            int zero = r - l - one;
            if (k <= zero) {
                l -= rank_l;
                r -= rank_r;
            } else {
                val |= (T)1 << bit;
                l = border[bit] + rank_l;
                r = border[bit] + rank_r;
                k -= zero;
            }
        }
        return val;
    }
};

}  // namespace ebi

#line 2 "data_structure/WaveletMatrix.hpp"

#line 2 "data_structure/bitVector.hpp"

#line 2 "template/int_alias.hpp"

#include <cstdint>

namespace ebi {

using ld = long double;
using std::size_t;
using i8 = std::int8_t;
using u8 = std::uint8_t;
using i16 = std::int16_t;
using u16 = std::uint16_t;
using i32 = std::int32_t;
using u32 = std::uint32_t;
using i64 = std::int64_t;
using u64 = std::uint64_t;
using i128 = __int128_t;
using u128 = __uint128_t;

}  // namespace ebi
#line 4 "data_structure/bitVector.hpp"

/*
    reference: https://misteer.hatenablog.com/entry/bit-vector
*/

#include <vector>


namespace ebi {

struct bitVector {
    u32 length, cn, bn;
    static u32 cw,
        bw;  // chunk, block の長さ cw = (lg N)^2, bw = (lg N)/2 とする.

    std::vector<u16> bit;
    std::vector<u32> chunk;
    std::vector<std::vector<u16>> blocks;

    bitVector(int n) : length(n) {
        cn = (length + cw - 1) / cw;
        bn = cw / bw;
        bit.assign(cn * bn, 0);
        chunk.assign(cn + 1, 0);
        blocks.assign(cn, std::vector<u16>(bn, 0));
    }

    void set(int k, int b) {
        u32 i = k / bw;
        u32 j = k % bw;
        if (b == 0) {
            bit[i] &= ~(1 << j);
        } else if (b == 1) {
            bit[i] |= (1 << j);
        }
    }

    int access(int k) {
        u32 i = k / bw;
        u32 j = k % bw;
        return bit[i] >> j & 1;
    }

    void build() {
        chunk[0] = 0;
        for (u32 i = 0; i < cn; i++) {
            blocks[i][0] = 0;
            for (u32 j = 0; j < bn - 1; j++) {
                blocks[i][j + 1] =
                    blocks[i][j] + __builtin_popcount(bit[i * bn + j]);
            }
            chunk[i + 1] = chunk[i] + blocks[i][bn - 1] +
                           __builtin_popcount(bit[(i + 1) * bn - 1]);
        }
    }

    // [0, r)に存在する1の数をO(1)で計算する.

    int rank(int r) {
        u32 i = r / cw;
        u32 j = (r % cw) / bw;
        u32 k = r % bw;
        u32 leftover = bit[i * bn + j] & ((1 << k) - 1);
        if (i == cn) return chunk[i];
        return chunk[i] + blocks[i][j] + __builtin_popcount(leftover);
    }

    int select(int k) {
        if (k == 0) return 0;
        if (rank(length) < k) return -1;
        u32 l = 0;
        u32 r = length;
        while (r - l > 1) {
            u32 mid = (r + l) / 2;
            if (rank(mid) >= k) {
                r = mid;
            } else {
                l = mid;
            }
        }
        return r;
    }

    int select0(int k) {
        if (k == 0) return 0;
        if (length - rank(length) < (u32)k) return -1;
        u32 l = 0;
        u32 r = length;
        while (r - l > 1) {
            u32 mid = (r + l) / 2;
            if (mid - rank(mid) >= (u32)k) {
                r = mid;
            } else {
                l = mid;
            }
        }
        return r;
    }
};

u32 bitVector::cw = 1024;
u32 bitVector::bw = 16;

}  // namespace ebi
#line 4 "data_structure/WaveletMatrix.hpp"

/*
    reference: https://miti-7.hatenablog.com/entry/2018/04/28/152259
*/

#include <limits>

#include <map>

#line 12 "data_structure/WaveletMatrix.hpp"

namespace ebi {

template <class T> struct WaveletMatrix {
  private:
    int wordsize = std::numeric_limits<T>::digits;
    int n;
    std::vector<bitVector> mat;
    std::vector<int> border;
    std::map<T, int> map;

  public:
    WaveletMatrix(std::vector<T> v)
        : n(v.size()), mat(wordsize, bitVector(n)), border(wordsize) {
        for (int bit = wordsize - 1; bit >= 0; --bit) {
            std::vector<T> one, zero;
            zero.reserve(n);
            for (int i = 0; i < n; ++i) {
                if (v[i] & ((T)1 << bit)) {
                    one.emplace_back(v[i]);
                    mat[bit].set(i, 1);
                } else {
                    zero.emplace_back(v[i]);
                }
            }
            border[bit] = zero.size();
            mat[bit].build();
            std::copy(one.begin(), one.end(), std::back_inserter(zero));
            v = zero;
        }
        for (int i = 0; i < n; i++) {
            if (map[v[i]] == 0) {
                map[v[i]] = i;
            }
        }
    }

    T access(int i) {
        T val = 0;
        int k = i;
        for (int bit = wordsize - 1; bit >= 0; --bit) {
            if (mat[bit].access(k) == 1) {
                val |= ((T)1 << bit);
                k = border[bit] + mat[bit].rank(k);
            } else {
                k = k - mat[bit].rank(k);
            }
        }
        return val;
    }

    int rank(int r, T x) {
        int k = r;
        for (int bit = wordsize - 1; bit >= 0; --bit) {
            if (x & ((T)1 << bit)) {
                k = border[bit] + mat[bit].rank(k);
            } else {
                k = k - mat[bit].rank(k);
            }
        }
        return k - map[x];
    }

    int select(T x, int k) {
        k = map[x] + k;
        for (int bit = 0; bit < wordsize; ++bit) {
            if (x & ((T)1 << bit)) {
                k = mat[bit].select(k - border[bit]);
            } else {
                k = mat[bit].select0(k);
            }
        }
        return k - 1;
    }

    T quantile(int l, int r, int k) {
        T val = 0;
        for (int bit = wordsize - 1; bit >= 0; --bit) {
            int rank_l = mat[bit].rank(l);
            int rank_r = mat[bit].rank(r);
            int one = rank_r - rank_l;
            int zero = r - l - one;
            if (k <= zero) {
                l -= rank_l;
                r -= rank_r;
            } else {
                val |= (T)1 << bit;
                l = border[bit] + rank_l;
                r = border[bit] + rank_r;
                k -= zero;
            }
        }
        return val;
    }
};

}  // namespace ebi

Back to top page