test/library-checker/DataStructure/RangeKthSmallest.test.cpp

• View this file on GitHub
• Last update: 2025-11-24 18:49:07+09:00
• Problem: https://judge.yosupo.jp/problem/range_kth_smallest

Depends on

• build/pch/stdc++.hpp
• library/datastructure/FullyIndexableDictionary.hpp
• library/datastructure/WaveletMatrix.hpp
• library/util/Compress.hpp

Code

#define PROBLEM "https://judge.yosupo.jp/problem/range_kth_smallest"
#include <bits/stdc++.h>

#include "library/datastructure/WaveletMatrix.hpp"

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    int n, q;
    std::cin >> n >> q;
    std::vector<int> v(n);
    for (int i = 0; i < n; i++)
        std::cin >> v[i];
    WaveletMatrix<int> WM(v);
    while (q--) {
        int l, r, k;
        std::cin >> l >> r >> k;
        std::cout << WM.kth_smallest(l, r, k) << "\n";
    }
}

#line 1 "test/library-checker/DataStructure/RangeKthSmallest.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/range_kth_smallest"
#include <bits/stdc++.h>

#line 4 "library/datastructure/FullyIndexableDictionary.hpp"
class FullyIndexableDictionary {
    int n,
        block; // 64個事に区切ったブロックの個数
    std::vector<unsigned long long> bit;
    std::vector<unsigned int> sum; // ブロック毎の累積和
    bool prepared;

  public:
    FullyIndexableDictionary() {}
    FullyIndexableDictionary(int n)
        : n(n), block((n + 63) >> 6), bit(block, 0), sum(block + 1, 0),
          prepared(false) {}

    bool is_prepared() { return prepared; }

    void set(int k) {
        bit[k >> 6] |= 1ull << (k & 63);
        sum[(k >> 6) + 1]++;
    }
    void build() {
        assert(!prepared);
        prepared = true;
        for (int i = 0; i < block; i++)
            sum[i + 1] += sum[i];
    }

    bool operator[](int k) const { return bool((bit[k >> 6] >> (k & 63)) & 1); }

    // [0,j) の合計
    int rank(int j, bool f = 1) {
        assert(prepared);
        int a = sum[j >> 6] +
                __builtin_popcountll(bit[j >> 6] & ((1ull << (j & 63)) - 1));
        return (f ? a : j - a);
    }
    // 0-indexed で k 番目の f の場所
    int select(int k, bool f = 1) {
        assert(prepared);
        if (k < 0 or rank(n, f) <= k)
            return -1;
        int l = 0, r = n;
        while (r - l > 1) {
            int m = (l + r) >> 1;
            (rank(m, f) >= k + 1 ? r : l) = m;
        }
        return r - 1;
    }
    // l以上で k 番目の f の場所
    int select(int k, bool f, int l) { return select(rank(l, f) + k, f); }
};
#line 2 "library/util/Compress.hpp"
template <typename T, bool Sentinel = false> class Compress {
    std::vector<T> v;
    bool prepared;

  public:
    Compress() : prepared(false) {
        if constexpr (Sentinel) {
            static_assert(std::numeric_limits<T>::is_specialized,
                          "cannot use Sentinel");
            v = {std::numeric_limits<T>::min(), std::numeric_limits<T>::max()};
        }
    }
    Compress(const std::vector<T> &w) : v(w), prepared(false) {
        if constexpr (Sentinel) {
            static_assert(std::numeric_limits<T>::is_specialized,
                          "cannot use Sentinel");
            v.push_back(std::numeric_limits<T>::min());
            v.push_back(std::numeric_limits<T>::max());
        }
        build();
    }

    void add(T a) {
        assert(!prepared);
        v.push_back(a);
    }
    void build() {
        assert(!prepared);
        prepared = true;
        std::ranges::sort(v);
        auto result = std::ranges::unique(v);
        v.erase(result.begin(), result.end());
    }

    bool is_prepared() const { return prepared; }

    int operator[](const T &a) const {
        assert(prepared);
        auto it = std::ranges::lower_bound(v, a);
        assert(*it == a);
        return std::distance(v.begin(), it);
    }
    int geq(const T &a) const {
        assert(prepared);
        auto it = std::ranges::lower_bound(v, a);
        return std::distance(v.begin(), it);
    }
    int gt(const T &a) const {
        assert(prepared);
        auto it = std::ranges::upper_bound(v, a);
        return std::distance(v.begin(), it);
    }
    int leq(const T &a) const {
        assert(prepared);
        auto it = --std::ranges::upper_bound(v, a);
        return std::distance(v.begin(), it);
    }
    int lt(const T &a) const {
        assert(prepared);
        auto it = --std::ranges::lower_bound(v, a);
        return std::distance(v.begin(), it);
    }
    T r(int id) const {
        assert(prepared);
        return v[id];
    }
    bool exist(const T &a) const {
        assert(prepared);
        return (*std::ranges::lower_bound(v, a)) == a;
    }
    int size() const { return v.size(); }
    T max() const { return v.back(); }
    T min() const { return v[0]; }

    friend std::ostream &operator<<(std::ostream &os, const Compress &C) {
        for (int i = 0; i < C.v.size(); i++)
            os << C.v[i] << ":" << i << " ";
        return os;
    }
};
#line 4 "library/datastructure/WaveletMatrix.hpp"
#define REP_(i, n) for (int i = 0; i < (n); i++)
template <typename T, bool COMPRESS = true> class WaveletMatrix {
  protected:
    using U = std::conditional_t<COMPRESS, int, T>;
    static_assert(std::is_integral_v<U>, "Wavelet Matrix is only for integer");
    int n, memo, log;
    std::vector<FullyIndexableDictionary> mat;
    std::vector<int> zero_cnt;
    Compress<T, true> C;
    std::vector<T> data;

    constexpr U comp(const T &x) const {
        if constexpr (COMPRESS) {
            return C.geq(x);
        } else {
            return x;
        }
    }
    constexpr T uncomp(const U &a) {
        if constexpr (COMPRESS) {
            return C.r(a);
        } else {
            return a;
        }
    }

    // 0-indexed で下から i bit 目
    inline bool low_bit(const U &a, int i) const { return (a >> i) & 1; }
    // 0-indexed で上から i bit 目
    inline bool high_bit(const U &a, int i) const {
        return low_bit(a, log - i - 1);
    }

    int nxt(int idx, int h, const U &a) {
        // idx の位置に a があった場合上から h bit 目でどこに行くか
        bool bit = high_bit(a, h);
        return mat[h].rank(idx, bit) + (bit ? zero_cnt[h] : 0);
    }

  public:
    WaveletMatrix(std::vector<T> v, int log_ = 0)
        : n(v.size()), data(v), log(log_) {
        std::vector<U> cv(n);
        if constexpr (COMPRESS) {
            C = Compress<T, true>(v);
            for (int i = 0; i < n; i++)
                cv[i] = C[v[i]];
            while (C.size() >= (1ull << log))
                log++;
        } else {
            cv = v;
            T mx = 0;
            for (const T &a : v) {
                assert(a >= 0);
                if (mx < a)
                    mx = a;
            }
            while (mx >= (1ull << log))
                log++;
        }
        mat.resize(log);
        zero_cnt.resize(log);
        std::vector<U> lv(n), rv(n);
        REP_(h, log) {
            mat[h] = FullyIndexableDictionary(n);
            int l = 0, r = 0;
            REP_(i, n)
            if (high_bit(cv[i], h)) {
                rv[r++] = cv[i];
                mat[h].set(i);
            } else
                lv[l++] = cv[i];
            zero_cnt[h] = l;
            mat[h].build();
            std::swap(lv, cv);
            REP_(i, r) cv[l + i] = rv[i];
        }
    }

    // [l,r) の x の個数
    int rank(int l, int r, const T &x) {
        if constexpr (COMPRESS) {
            if (!C.exist(x))
                return 0;
        }
        U a = comp(x);
        REP_(h, log) {
            l = nxt(l, h, a);
            r = nxt(r, h, a);
        }
        memo = l;
        return r - l;
    }
    int rank(int r, const T &x) { return rank(x, 0, r); }

    // k 番目の x の index
    int select(int k, const T &x) {
        U a = comp(x);
        if (rank(a, n) < k)
            return -1;
        k += memo;
        for (int h = log - 1; h >= 0; h--) {
            bool bit = high_bit(a, h);
            if (bit)
                k -= zero_cnt[h];
            k = mat[h].select(k, bit);
        }
        return k;
    }

    // [l,r) で 0-indexed k 番目に大きい値
    T kth_largest(int l, int r, int k) {
        if (k < 0 or r - l <= k)
            return -1;
        U res = 0;
        REP_(h, log) {
            int L = mat[h].rank(l);
            int R = mat[h].rank(r);
            res <<= 1;
            if (R - L > k) {
                l = L + zero_cnt[h];
                r = R + zero_cnt[h];
                res++;
            } else {
                k -= R - L;
                l -= L;
                r -= R;
            }
        }
        return uncomp(res);
    }
    T kth_smallest(int l, int r, int k) {
        return kth_largest(l, r, r - l - k - 1);
    }

    // [l,r) で x 未満の個数
    int range_freq(int l, int r, const T &upper) {
        U a = comp(upper);
        int res = 0;
        REP_(h, log) {
            if (high_bit(a, h)) {
                int L = mat[h].rank(l, 0), R = mat[h].rank(r, 0);
                res += R - L;
            }
            l = nxt(l, h, a);
            r = nxt(r, h, a);
        }
        return res;
    }
    // [l,r) で [lower,upper) の個数
    int range_freq(int l, int r, const T &lower, const T &upper) {
        return range_freq(l, r, upper) - range_freq(l, r, lower);
    }

    std::optional<T> lt(int l, int r, const T &x) {
        int cnt = range_freq(l, r, x);
        if (cnt)
            return kth_smallest(l, r, cnt - 1);
        return std::nullopt;
    }
    std::optional<T> leq(int l, int r, const T &x) {
        if (rank(l, r, x))
            return x;
        return lt(l, r, x);
    }
    std::optional<T> geq(int l, int r, const T &x) {
        int cnt = r - l - range_freq(l, r, x);
        if (cnt)
            return kth_largest(l, r, cnt - 1);
        return std::nullopt;
    }
    std::optional<T> gt(int l, int r, const T &x) {
        T y;
        if constexpr (COMPRESS) {
            y = C.r(C.gt(x));
        } else {
            y = x + 1;
        }
        return geq(l, r, y);
    }

    // セグ木などを載せる時用
    // BIT は専用のライブラリを作ってあるが、一応抽象化も持っておく
    // 構築したものを返してるので遅そう
    int height() const { return log; }
    std::vector<int> points(int idx) {
        std::vector<int> res(log);
        U a = comp(data[idx]);
        REP_(h, log) {
            idx = nxt(idx, h, a);
            res[h] = idx;
        }
        return res;
    }
    std::vector<std::tuple<int, int, int>> intervals(int l, int r,
                                                     const T &upper) {
        std::vector<std::tuple<int, int, int>> res;
        U a = comp(upper);
        REP_(h, log) {
            if (high_bit(a, h)) {
                int L = mat[h].rank(l, 0), R = mat[h].rank(r, 0);
                res.emplace_back(h, L, R);
            }
            l = nxt(l, h, a);
            r = nxt(r, h, a);
        }
        return res;
    }
    std::vector<std::tuple<int, int, int>> kth_largest_intervals(int l, int r,
                                                                 int k) {
        assert(0 <= k and k < r - l);
        std::vector<std::tuple<int, int, int>> res;
        REP_(h, log) {
            int L = mat[h].rank(l);
            int R = mat[h].rank(r);
            if (R - L > k) {
                l = L + zero_cnt[h];
                r = R + zero_cnt[h];
            } else {
                res.emplace_back(h, L + zero_cnt[h], R + zero_cnt[h]);
                k -= R - L;
                l -= L;
                r -= R;
            }
        }
        return res;
    }
    std::vector<std::tuple<int, int, int>> kth_smallest_intervals(int l, int r,
                                                                  int k) {
        return kth_largest_intervals(l, r, r - l - k - 1);
    }
};
#undef REP_
#line 5 "test/library-checker/DataStructure/RangeKthSmallest.test.cpp"

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    int n, q;
    std::cin >> n >> q;
    std::vector<int> v(n);
    for (int i = 0; i < n; i++)
        std::cin >> v[i];
    WaveletMatrix<int> WM(v);
    while (q--) {
        int l, r, k;
        std::cin >> l >> r >> k;
        std::cout << WM.kth_smallest(l, r, k) << "\n";
    }
}

Back to top page