test/yosupo/range_kth_smallest.test.cpp

View this file on GitHub
Last update: 2024-01-08 02:22:28+09:00
Problem: https://judge.yosupo.jp/problem/range_kth_smallest

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/range_kth_smallest"

#include <bits/stdc++.h>

#include "../../data-structure/wavelet_matrix.hpp"

using namespace std;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0);

    int N, Q;
    cin >> N >> Q;
    vector<int> a(N);
    for (int i = 0; i < N; ++i) cin >> a[i];
    WaveletMatrix<int, 30> wm(a);
    for (int i = 0; i < Q; ++i) {
        int l, r, k;
        cin >> l >> r >> k;
        cout << wm.quantile(l, r, k) << "\n";
    }
}

#line 1 "test/yosupo/range_kth_smallest.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/range_kth_smallest"

#include <bits/stdc++.h>

#line 5 "data-structure/wavelet_matrix.hpp"

#line 2 "data-structure/bit_vector.hpp"
#include <bit>
#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;
};
#line 6 "test/yosupo/range_kth_smallest.test.cpp"

using namespace std;

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0);

    int N, Q;
    cin >> N >> Q;
    vector<int> a(N);
    for (int i = 0; i < N; ++i) cin >> a[i];
    WaveletMatrix<int, 30> wm(a);
    for (int i = 0; i < Q; ++i) {
        int l, r, k;
        cin >> l >> r >> k;
        cout << wm.quantile(l, r, k) << "\n";
    }
}