Test/LC/range_kth_smallest.test.cpp

View this file on GitHub
Last update: 2025-03-04 23:23:46+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 "../../Src/Template/IOSetting.hpp"
#include "../../Src/Sequence/RangeKthSmallest.hpp"

#include <iostream>
#include <vector>

using namespace zawa;

struct query {
    usize k, l, r;
};

int main() {
    SetFastIO();
    int N, Q;
    std::cin >> N >> Q;
    std::vector<int> A(N);
    for (auto& a : A) std::cin >> a;
    std::vector<query> q(Q);
    for (int i{} ; i < Q ; i++) {
        usize l, r, k;
        std::cin >> l >> r >> k;
        q[i] = { k, l, r };
    }
    for (auto ans : RangeKthSmallest(A, q)) {
        std::cout << ans << '\n';
    }
}

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

#line 2 "Src/Template/IOSetting.hpp"

#line 2 "Src/Template/TypeAlias.hpp"

#include <cstdint>
#include <cstddef>

namespace zawa {

using i16 = std::int16_t;
using i32 = std::int32_t;
using i64 = std::int64_t;
using i128 = __int128_t;

using u8 = std::uint8_t;
using u16 = std::uint16_t;
using u32 = std::uint32_t;
using u64 = std::uint64_t;

using usize = std::size_t;

} // namespace zawa
#line 4 "Src/Template/IOSetting.hpp"

#include <iostream>
#include <iomanip>

namespace zawa {

void SetFastIO() {
    std::cin.tie(nullptr)->sync_with_stdio(false);
}

void SetPrecision(u32 dig) {
    std::cout << std::fixed << std::setprecision(dig);
}

} // namespace zawa
#line 2 "Src/Sequence/RangeKthSmallest.hpp"

#line 2 "Src/Sequence/CompressedSequence.hpp"

#line 4 "Src/Sequence/CompressedSequence.hpp"

#include <vector>
#include <algorithm>
#include <cassert>
#include <iterator>
#include <limits>

namespace zawa {

template <class T>
class CompressedSequence {
public:

    static constexpr u32 NotFound = std::numeric_limits<u32>::max();

    CompressedSequence() = default;

    template <class InputIterator>
    CompressedSequence(InputIterator first, InputIterator last) : comped_(first, last), f_{} {
        std::sort(comped_.begin(), comped_.end());
        comped_.erase(std::unique(comped_.begin(), comped_.end()), comped_.end());
        comped_.shrink_to_fit();
        f_.reserve(std::distance(first, last));
        for (auto it{first} ; it != last ; it++) {
            f_.emplace_back(std::distance(comped_.begin(), std::lower_bound(comped_.begin(), comped_.end(), *it)));
        }
    }

    CompressedSequence(const std::vector<T>& A) : CompressedSequence(A.begin(), A.end()) {}

    inline usize size() const noexcept {
        return comped_.size();
    }

    u32 operator[](const T& v) const {
        return std::distance(comped_.begin(), std::lower_bound(comped_.begin(), comped_.end(), v));
    }

    u32 upper_bound(const T& v) const {
        return std::distance(comped_.begin(), std::upper_bound(comped_.begin(), comped_.end(), v));
    }

    u32 find(const T& v) const {
        u32 i = std::distance(comped_.begin(), std::lower_bound(comped_.begin(), comped_.end(), v));
        return i == comped_.size() or comped_[i] != v ? NotFound : i;
    }

    bool contains(const T& v) const {
        u32 i = std::distance(comped_.begin(), std::lower_bound(comped_.begin(), comped_.end(), v));
        return i < comped_.size() and comped_[i] == v;
    }

    u32 at(const T& v) const {
        u32 res = find(v);
        assert(res != NotFound);
        return res;
    }

    inline u32 map(u32 i) const noexcept {
        assert(i < f_.size());
        return f_[i];
    }

    inline T inverse(u32 i) const noexcept {
        assert(i < size());
        return comped_[i];
    }

    inline std::vector<T> comped() const noexcept {
        return comped_;
    }

private:

    std::vector<T> comped_;

    std::vector<u32> f_;

};

} // namespace zawa
#line 5 "Src/Sequence/RangeKthSmallest.hpp"

#line 7 "Src/Sequence/RangeKthSmallest.hpp"
#include <bit>
#line 9 "Src/Sequence/RangeKthSmallest.hpp"
#include <utility>
#line 11 "Src/Sequence/RangeKthSmallest.hpp"

namespace zawa {

template <class T, class U>
std::vector<T> RangeKthSmallest(const std::vector<T>& A, const std::vector<U>& Q) {
    assert(A.size());
    CompressedSequence comp{A};
    std::vector<std::vector<std::pair<usize, bool>>> event(A.size() + 1);
    for (usize i{} ; i < Q.size() ; i++) {
        const U& q{Q[i]};
        assert(q.l < q.r and q.r <= A.size());
        assert(q.k < q.r - q.l);
        event[q.l].emplace_back(i, true);
        event[q.r].emplace_back(i, false);
    }
    std::vector<usize> x(Q.size()), cnt(Q.size());
    const usize LOG{(usize)(std::bit_width(comp.size()) -
            (std::has_single_bit(comp.size()) ? 1 : 0))};
    for (usize i{LOG} ; i-- ; ) {
        std::vector<i32> seg(1 << (LOG - i));
        std::vector<i32> add(Q.size());
        for (usize j{1} ; j <= A.size() ; j++) {
            seg[comp.map(j - 1) >> i]++;
            for (auto [k, s] : event[j]) {
                add[k] += seg[(x[k] >> i)] * (s ? -1 : 1);
            }
        }
        for (u32 j{} ; j < Q.size() ; j++) {
            if (cnt[j] + add[j] <= Q[j].k) {
                x[j] += 1 << i;
                cnt[j] += add[j];
            }
        }
    } 
    std::vector<T> res(Q.size());
    for (usize i{} ; i < Q.size() ; i++) {
        res[i] = comp.inverse(x[i]);
    }
    return res;
}

} // namespace zawa
#line 5 "Test/LC/range_kth_smallest.test.cpp"

#line 8 "Test/LC/range_kth_smallest.test.cpp"

using namespace zawa;

struct query {
    usize k, l, r;
};

int main() {
    SetFastIO();
    int N, Q;
    std::cin >> N >> Q;
    std::vector<int> A(N);
    for (auto& a : A) std::cin >> a;
    std::vector<query> q(Q);
    for (int i{} ; i < Q ; i++) {
        usize l, r, k;
        std::cin >> l >> r >> k;
        q[i] = { k, l, r };
    }
    for (auto ans : RangeKthSmallest(A, q)) {
        std::cout << ans << '\n';
    }
}