data_structure/kth_element.hpp

View this file on GitHub
Last update: 2024-05-13 20:26:20+09:00
Include: #include "data_structure/kth_element.hpp"
Link: https://qiita.com/drken/items/1b7e6e459c24a83bb7fd

Depends on

Required by

Verified with

Code

#pragma once


#include <cassert>
#include <vector>
#include <functional>
#include <queue>
#include <optional>
#include <concepts>
#include <ranges>


#include "internal/dev_env.hpp"
#include "internal/types.hpp"


namespace uni {


// Thanks to: https://qiita.com/drken/items/1b7e6e459c24a83bb7fd
template<
    class T,
    std::ranges::output_range<T> Container = std::vector<T>,
    std::strict_weak_order<T, T> Comparer = std::less<T>,
    std::strict_weak_order<T, T> RevComparer = std::greater<T>,
    std::integral Size = internal::size_t
>
struct kth_element {
    using value_type = T;
    using size_type = Size;

  protected:
    size_type _k;
    std::priority_queue<T, Container, Comparer> _small;
    std::priority_queue<T, Container, RevComparer> _large;

  public:
    kth_element(const size_type k = 0) noexcept(NO_EXCEPT) : _k(k + 1) { assert(k >= 0); }

    inline bool has() const noexcept(NO_EXCEPT) { return std::ssize(this->_small) == this->_k; }

    inline value_type value() const noexcept(NO_EXCEPT) { return this->_small.top(); }

    inline std::optional<value_type> get() const noexcept(NO_EXCEPT) {
        if(this->has()) return this->value();
        return {};
    }

    template<std::convertible_to<T> U = T>
    inline auto value_or(U&& v) const noexcept(NO_EXCEPT) {
        return this->get().value_or(std::forward<U>(v));
    }

    inline void push(const value_type& v) noexcept(NO_EXCEPT) {
        if(std::ssize(this->_small) < this->_k) {
            this->_small.push(v);
            return;
        }

        const auto kth = this->_small.top();

        if(Comparer{}(v, kth)) {
            this->_small.pop();
            this->_small.push(v);
            this->_large.push(kth);
        }
        else {
            this->_large.push(v);
        }
    }

    inline void pop() noexcept(NO_EXCEPT) {
        assert(this->has());

        this->_small.pop();

        if(this->_large.empty()) return;

        const auto v = this->_large.top();

        this->_large.pop();
        this->_small.push(v);
    }
};


} // namespace uni

#line 2 "data_structure/kth_element.hpp"


#include <cassert>
#include <vector>
#include <functional>
#include <queue>
#include <optional>
#include <concepts>
#include <ranges>


#line 2 "internal/dev_env.hpp"


#ifdef LOCAL_JUDGE
    inline constexpr bool DEV_ENV = true;
    inline constexpr bool NO_EXCEPT = false;
#else
    inline constexpr bool DEV_ENV = false;
    inline constexpr bool NO_EXCEPT = true;
#endif // LOCAL_JUDGE


#if __cplusplus >= 202100L
    #define CPP20 true
    #define CPP23 true
#elif __cplusplus >= 202002L
    #define CPP20 true
    #define CPP23 false
#else
    #define CPP20 false
    #define CPP23 false
#endif
#line 2 "internal/types.hpp"

#include <cstdint>

namespace uni {

namespace internal {


using size_t = std::int64_t;

using int128_t = __int128_t;
using uint128_t = __uint128_t;


} // namesapce internal

} // namespace uni
#line 15 "data_structure/kth_element.hpp"


namespace uni {


// Thanks to: https://qiita.com/drken/items/1b7e6e459c24a83bb7fd
template<
    class T,
    std::ranges::output_range<T> Container = std::vector<T>,
    std::strict_weak_order<T, T> Comparer = std::less<T>,
    std::strict_weak_order<T, T> RevComparer = std::greater<T>,
    std::integral Size = internal::size_t
>
struct kth_element {
    using value_type = T;
    using size_type = Size;

  protected:
    size_type _k;
    std::priority_queue<T, Container, Comparer> _small;
    std::priority_queue<T, Container, RevComparer> _large;

  public:
    kth_element(const size_type k = 0) noexcept(NO_EXCEPT) : _k(k + 1) { assert(k >= 0); }

    inline bool has() const noexcept(NO_EXCEPT) { return std::ssize(this->_small) == this->_k; }

    inline value_type value() const noexcept(NO_EXCEPT) { return this->_small.top(); }

    inline std::optional<value_type> get() const noexcept(NO_EXCEPT) {
        if(this->has()) return this->value();
        return {};
    }

    template<std::convertible_to<T> U = T>
    inline auto value_or(U&& v) const noexcept(NO_EXCEPT) {
        return this->get().value_or(std::forward<U>(v));
    }

    inline void push(const value_type& v) noexcept(NO_EXCEPT) {
        if(std::ssize(this->_small) < this->_k) {
            this->_small.push(v);
            return;
        }

        const auto kth = this->_small.top();

        if(Comparer{}(v, kth)) {
            this->_small.pop();
            this->_small.push(v);
            this->_large.push(kth);
        }
        else {
            this->_large.push(v);
        }
    }

    inline void pop() noexcept(NO_EXCEPT) {
        assert(this->has());

        this->_small.pop();

        if(this->_large.empty()) return;

        const auto v = this->_large.top();

        this->_large.pop();
        this->_small.push(v);
    }
};


} // namespace uni