test/yosupo/ds/range_kth_smallest.test.cpp

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

Depends on

• Wavelet Matrix (ds/wavelet_matrix.hpp)

Code

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

#include <bits/stdc++.h>
using namespace std;

#include "ds/wavelet_matrix.hpp"

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    int n, q;
    cin >> n >> q;
    vector<int> v(n);
    int mx = 0;
    for (int &i : v) cin >> i, mx = max(mx, i);
    WaveletMatrix<int> wm(v);
    while (q--) {
        int l, r, k;
        cin >> l >> r >> k;
        cout << wm.kth_smallest(l, r, k) << "\n";
    }
}

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

#include <bits/stdc++.h>
using namespace std;

#line 2 "ds/wavelet_matrix.hpp"

using u32 = uint32_t;
using u64 = uint64_t;

// clang-format off
struct BitVec {
    BitVec() = default;
    explicit BitVec(int n, const vector<bool> &a) :
        BitVec(n, [&](int i) { return a[i]; }) {}
    template<class Gen> BitVec(int n, Gen gen) : v(n / wordsize + 1) {
        for (int i = 0; i < n; i++)
            v[i / wordsize].bit |= size_t(gen(i)) << i % wordsize;
        for (int i = 1; i < v.size(); i++)
            v[i].sum = v[i - 1].sum + popcnt64(v[i - 1].bit);
    }
    bool operator[](int i) const {
        return v[i / wordsize].bit >> i % wordsize & 1;
    }
    int rank(bool val, int i) const {
        int x = v[i / wordsize].sum +
                popcnt64(v[i / wordsize].bit & ~(~size_t(0) << i % wordsize));
        return val ? x : i - x;
    }
    int rank(bool val, int l, int r) const {
        return rank(val, r) - rank(val, l);
    }

  private:
    static constexpr int wordsize = numeric_limits<size_t>::digits;
    struct node { size_t bit = 0, sum = 0; };
    vector<node> v;
    static inline constexpr int popcnt64(size_t x) {
#ifdef __GNUC__
        return __builtin_popcountll(x);
#endif
        x -= x >> 1 & 0x5555555555555555;
        x = (x & 0x3333333333333333) + (x >> 2 & 0x3333333333333333);
        x = (x + (x >> 4)) & 0x0f0f0f0f0f0f0f0f;
        return x * 0x0101010101010101 >> 56;
    }
};

// suisen-cp's cp-library-cpp
template<class T, int bit_num = numeric_limits<T>::digits>
struct WaveletMatrix {
    WaveletMatrix() = default;
    explicit WaveletMatrix(const vector<T> &a) :
        WaveletMatrix((int)a.size(), [&](int i) { return a[i]; }) {}
    template<class Gen> WaveletMatrix(int _n, Gen gen) : n(_n) {
        vector<T> a(n), l(n), r(n);
        for (int i = 0; i < n; i++) a[i] = gen(i);
        for (int lg = bit_num - 1; lg >= 0; lg--) {
            bv[lg] = BitVec(n, [&](int i) { return a[i] >> lg & 1; });
            int li = 0, ri = 0;
            for (int i = 0; i < n; i++)
                (a[i] >> lg & 1 ? r[ri++] : l[li++]) = a[i];
            a.swap(l);
            copy(begin(r), begin(r) + ri, begin(a) + li);
            mid[lg] = li;
        }
    }
    T operator[](int i) const {
        T res = 0;
        for (int lg = bit_num - 1; lg >= 0; lg--) {
            bool b = bv[lg][i];
            res |= T(b) << lg;
            i = b * mid[lg] + bv[lg].rank(b, i);
        }
        return res;
    }
    T rank(T val, int i) const { return rank(val, 0, i); }
    T rank(T val, int l, int r) const {
        for (int lg = bit_num - 1; lg >= 0; lg--)
            succ(l, r, val >> lg & 1, lg);
        return r - l;
    }
    T kth_smallest(int l, int r, int k) const {
        T res = 0;
        for (int lg = bit_num - 1; lg >= 0; lg--) {
            int cnt = bv[lg].rank(0, l, r), bit = k >= cnt;
            succ(l, r, bit, lg);
            res |= T(bit) << lg;
            k -= bit * cnt;
        }
        return res;
    }
    int freq(int l, int r, T upper) const {
        int res = 0;
        for (int log = bit_num - 1; log >= 0; --log) {
            bool u = upper >> log & 1;
            if (u) res += bv[log].rank(0, l, r);
            succ(l, r, u, log);
        }
        return res;
    }
    int freq(int l, int r, T lower, T upper) const {
        return freq(l, r, upper) - freq(l, r, lower);
    }

  private:
    int n = 0;
    array<BitVec, bit_num> bv;
    array<int, bit_num> mid;
    void succ(int &l, int &r, bool b, int lg) const {
        l = b * mid[lg] + bv[lg].rank(b, l);
        r = b * mid[lg] + bv[lg].rank(b, r);
    }
};
// clang-format on
#line 7 "test/yosupo/ds/range_kth_smallest.test.cpp"

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    int n, q;
    cin >> n >> q;
    vector<int> v(n);
    int mx = 0;
    for (int &i : v) cin >> i, mx = max(mx, i);
    WaveletMatrix<int> wm(v);
    while (q--) {
        int l, r, k;
        cin >> l >> r >> k;
        cout << wm.kth_smallest(l, r, k) << "\n";
    }
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	all_zero_00	AC	6 ms	4 MB
g++	dense_large_a_00	AC	110 ms	4 MB
g++	dense_small_a_00	AC	109 ms	4 MB
g++	example_00	AC	5 ms	4 MB
g++	max_random_00	AC	276 ms	8 MB
g++	max_random_01	AC	289 ms	8 MB
g++	max_random_02	AC	279 ms	8 MB
g++	max_random_03	AC	275 ms	8 MB
g++	max_random_04	AC	275 ms	8 MB
g++	random_00	AC	194 ms	6 MB
g++	random_01	AC	199 ms	7 MB
g++	random_02	AC	121 ms	4 MB
g++	random_03	AC	71 ms	7 MB
g++	random_04	AC	78 ms	4 MB
g++	small_00	AC	5 ms	4 MB
g++	small_01	AC	4 ms	4 MB
g++	small_02	AC	4 ms	4 MB
g++	small_03	AC	4 ms	4 MB
g++	small_04	AC	5 ms	4 MB
g++	small_05	AC	4 ms	4 MB
g++	small_06	AC	4 ms	4 MB
g++	small_07	AC	4 ms	4 MB
g++	small_08	AC	4 ms	4 MB
g++	small_09	AC	5 ms	4 MB
clang++	all_zero_00	AC	6 ms	4 MB
clang++	dense_large_a_00	AC	99 ms	4 MB
clang++	dense_small_a_00	AC	99 ms	4 MB
clang++	example_00	AC	5 ms	4 MB
clang++	max_random_00	AC	231 ms	8 MB
clang++	max_random_01	AC	229 ms	8 MB
clang++	max_random_02	AC	230 ms	8 MB
clang++	max_random_03	AC	231 ms	8 MB
clang++	max_random_04	AC	231 ms	8 MB
clang++	random_00	AC	160 ms	6 MB
clang++	random_01	AC	170 ms	7 MB
clang++	random_02	AC	103 ms	4 MB
clang++	random_03	AC	54 ms	7 MB
clang++	random_04	AC	64 ms	4 MB
clang++	small_00	AC	5 ms	4 MB
clang++	small_01	AC	4 ms	4 MB
clang++	small_02	AC	4 ms	4 MB
clang++	small_03	AC	4 ms	4 MB
clang++	small_04	AC	4 ms	4 MB
clang++	small_05	AC	4 ms	4 MB
clang++	small_06	AC	4 ms	4 MB
clang++	small_07	AC	4 ms	4 MB
clang++	small_08	AC	4 ms	4 MB
clang++	small_09	AC	4 ms	4 MB

Back to top page