ウェーブレット行列
(data_structure/sequence/wavelet_matrix.hpp)

• View this file on GitHub
• Last update: 2022-07-19 22:00:44+09:00
• Include: #include "data_structure/sequence/wavelet_matrix.hpp"

Depends on

• コンパクト Bit Vector (data_structure/sequence/compact_bit_vector.hpp)

Verified with

• test/aoj/itp/itp1_6_a.test.cpp
• test/library_checker/data_structure/range_kth_smallest.test.cpp
• test/library_checker/data_structure/static_range_frequency.test.cpp

Code

#pragma once
/**
 * @brief ウェーブレット行列
*/
#include "compact_bit_vector.hpp"
struct wavelet_matrix{
  int N, LOG;
  vector<bit_vector> D;
  vector<int> cnt, T;
  wavelet_matrix(){
  }
  wavelet_matrix(vector<int> &A): N(A.size()){
    LOG = 1;
    for (int i = 0; i < N; i++){
      if (A[i] > 0){
        LOG = max(LOG, 32 - __builtin_clz(A[i]));
      }
    }
    D.resize(LOG);
    cnt.resize(LOG, 0);
    for (int i = LOG - 1; i >= 0; i--){
      vector<bool> B(N, false);
      for (int j = 0; j < N; j++){
        if ((A[j] >> i & 1) == 1){
          B[j] = true;
        }
      }
      D[LOG - 1 - i] = bit_vector(B);
      vector<int> A2;
      for (int j = 0; j < N; j++){
        if ((A[j] >> i & 1) == 0){
          A2.push_back(A[j]);
          cnt[LOG - 1 - i]++;
        }
      }
      for (int j = 0; j < N; j++){
        if ((A[j] >> i & 1) == 1){
          A2.push_back(A[j]);
        }
      }
      swap(A, A2);
    }
    T = A;
  }
  int operator [](int k){
    int ans = 0;
    for (int i = 0; i < LOG; i++){
      if (D[i][k] == 0){
        k = D[i].rank0(k);
      } else {
        ans += 1 << (LOG - 1 - i);
        k = cnt[i] + D[i].rank1(k);
      }
    }
    return ans;
  }
  int rank(int l, int r, int c){
    for (int i = 0; i < LOG; i++){
      if ((c >> (LOG - 1 - i) & 1) == 0){
        l = D[i].rank0(l);
        r = D[i].rank0(r);
      } else {
        l = cnt[i] + D[i].rank1(l);
        r = cnt[i] + D[i].rank1(r);
      }
    }
    return r - l;
  }
  int quantile(int l, int r, int k){
    int ans = 0;
    for (int i = 0; i < LOG; i++){
      int cnt0 = D[i].rank0(r) - D[i].rank0(l);
      if (k < cnt0){
        l = D[i].rank0(l);
        r = D[i].rank0(r);
      } else {
        ans += 1 << (LOG - 1 - i);
        k -= cnt0;
        l = cnt[i] + D[i].rank1(l);
        r = cnt[i] + D[i].rank1(r);
      }
    }
    return ans;
  }
};

#line 2 "data_structure/sequence/wavelet_matrix.hpp"
/**
 * @brief ウェーブレット行列
*/
#line 2 "data_structure/sequence/compact_bit_vector.hpp"
/**
 * @brief コンパクト Bit Vector
*/
struct bit_vector{
  vector<unsigned long long> A;
  vector<int> S;
  bit_vector(){
  }
  bit_vector(vector<bool> &a){
    int N = a.size();
    int M = (N + 64 - 1) >> 6;
    A = vector<unsigned long long>(M, 0);
    for (int i = 0; i < M; i++){
      for (int j = i << 6; j < min((i + 1) << 6, N); j++){
        if (a[j]){
          A[i] |= (unsigned long long) 1 << (j - (i << 6));
        }
      }
    }
    S = vector<int>(M + 1, 0);
    for (int i = 0; i < M; i++){
      S[i + 1] = S[i] + __builtin_popcountll(A[i]);
    }
  }
  int operator [](int k){
    return A[k >> 6] >> (k & 63) & 1;
  }
  int rank0(int k){
    return k - rank1(k);
  }
  int rank1(int k){
    if ((k & 63) == 0){
      return S[k >> 6];
    } else {
      return S[k >> 6] + __builtin_popcountll(A[k >> 6] << (64 - k + (k >> 6 << 6)));
    }
  }
};
#line 6 "data_structure/sequence/wavelet_matrix.hpp"
struct wavelet_matrix{
  int N, LOG;
  vector<bit_vector> D;
  vector<int> cnt, T;
  wavelet_matrix(){
  }
  wavelet_matrix(vector<int> &A): N(A.size()){
    LOG = 1;
    for (int i = 0; i < N; i++){
      if (A[i] > 0){
        LOG = max(LOG, 32 - __builtin_clz(A[i]));
      }
    }
    D.resize(LOG);
    cnt.resize(LOG, 0);
    for (int i = LOG - 1; i >= 0; i--){
      vector<bool> B(N, false);
      for (int j = 0; j < N; j++){
        if ((A[j] >> i & 1) == 1){
          B[j] = true;
        }
      }
      D[LOG - 1 - i] = bit_vector(B);
      vector<int> A2;
      for (int j = 0; j < N; j++){
        if ((A[j] >> i & 1) == 0){
          A2.push_back(A[j]);
          cnt[LOG - 1 - i]++;
        }
      }
      for (int j = 0; j < N; j++){
        if ((A[j] >> i & 1) == 1){
          A2.push_back(A[j]);
        }
      }
      swap(A, A2);
    }
    T = A;
  }
  int operator [](int k){
    int ans = 0;
    for (int i = 0; i < LOG; i++){
      if (D[i][k] == 0){
        k = D[i].rank0(k);
      } else {
        ans += 1 << (LOG - 1 - i);
        k = cnt[i] + D[i].rank1(k);
      }
    }
    return ans;
  }
  int rank(int l, int r, int c){
    for (int i = 0; i < LOG; i++){
      if ((c >> (LOG - 1 - i) & 1) == 0){
        l = D[i].rank0(l);
        r = D[i].rank0(r);
      } else {
        l = cnt[i] + D[i].rank1(l);
        r = cnt[i] + D[i].rank1(r);
      }
    }
    return r - l;
  }
  int quantile(int l, int r, int k){
    int ans = 0;
    for (int i = 0; i < LOG; i++){
      int cnt0 = D[i].rank0(r) - D[i].rank0(l);
      if (k < cnt0){
        l = D[i].rank0(l);
        r = D[i].rank0(r);
      } else {
        ans += 1 << (LOG - 1 - i);
        k -= cnt0;
        l = cnt[i] + D[i].rank1(l);
        r = cnt[i] + D[i].rank1(r);
      }
    }
    return ans;
  }
};

Back to top page