test/src/datastructure/double_ended_priority_queue/double_ended_priority_queue.test.cpp

• View this file on GitHub
• Last update: 2022-10-13 20:07:57+09:00
• Problem: https://judge.yosupo.jp/problem/double_ended_priority_queue

Depends on

• Double Ended Priority Queue (library/datastructure/double_ended_priority_queue.hpp)

Code

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

#include <iostream>

#include "library/datastructure/double_ended_priority_queue.hpp"

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

    int n, q;
    std::cin >> n >> q;

    std::vector<int> init(n);
    for (auto &e : init) std::cin >> e;
    suisen::DoubleEndedPriorityQueue<int> pq(init.begin(), init.end());

    for (int i = 0; i < q; ++i) {
        int query_type;
        std::cin >> query_type;
        if (query_type == 0) {
            int x;
            std::cin >> x;
            pq.push(x);
        } else if (query_type == 1) {
            std::cout << pq.pop_min() << '\n';
        } else {
            std::cout << pq.pop_max() << '\n';
        }
    }
    return 0;
}

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

#include <iostream>

#line 1 "library/datastructure/double_ended_priority_queue.hpp"



#include <algorithm>
#include <cassert>
#include <vector>
#include <functional>
#include <utility>

namespace suisen {
    template <typename T, typename Comp = std::less<T>>
    struct DoubleEndedPriorityQueue {
        using value_type = T;
        using compare_type = Comp;

        DoubleEndedPriorityQueue() = default;
        DoubleEndedPriorityQueue(const Comp& comp) : _comp(comp) {}

        template <typename InputIterator>
        DoubleEndedPriorityQueue(InputIterator first, InputIterator last) : _max_heap(first, last) {
            std::make_heap(_max_heap.begin(), _max_heap.end(), _comp);
        }
        template <typename InputIterator>
        DoubleEndedPriorityQueue(InputIterator first, InputIterator last, const Comp& comp) : _comp(comp), _max_heap(first, last) {
            std::make_heap(_max_heap.begin(), _max_heap.end(), _comp);
        }

        template <typename Iterable, typename = std::void_t<typename Iterable::iterator>>
        DoubleEndedPriorityQueue(const Iterable& dat) : DoubleEndedPriorityQueue(dat.begin(), dat.end()) {}
        template <typename Iterable, typename = std::void_t<typename Iterable::iterator>>
        DoubleEndedPriorityQueue(const Iterable& dat, Comp& comp) : DoubleEndedPriorityQueue(dat.begin(), dat.end(), comp) {}

        bool empty() const { return size() == 0; }
        int size() const { return _min_heap.size() + _max_heap.size(); }

        void push(const value_type& v) {
            if (_min_heap.empty() or _comp(pivot, v)) {
                _max_heap.push_back(v);
                std::push_heap(_max_heap.begin(), _max_heap.end(), _comp);
            } else {
                _min_heap.push_back(v);
                std::push_heap(_min_heap.begin(), _min_heap.end(), _rev_comp);
            }
        }
        template <typename ...Args>
        void emplace(Args &&...args) { push(value_type(std::forward<Args>(args)...)); }

        const value_type& max() const {
            assert(size());
            return _max_heap.size() ? _max_heap.front() : pivot;
        }
        const value_type& min() {
            ensure_min_heap_nonempty();
            return _min_heap.front();
        }
        const value_type& top() const { return max(); }

        value_type pop_max() {
            ensure_max_heap_nonempty();
            std::pop_heap(_max_heap.begin(), _max_heap.end(), _comp);
            value_type res = std::move(_max_heap.back());
            _max_heap.pop_back();
            return res;
        }
        value_type pop_min() {
            ensure_min_heap_nonempty();
            std::pop_heap(_min_heap.begin(), _min_heap.end(), _rev_comp);
            value_type res = std::move(_min_heap.back());
            _min_heap.pop_back();
            return res;
        }
        value_type pop() { return pop_max(); }

        std::vector<value_type> dump_sorted() const {
            std::vector<value_type> res_l(_min_heap), res_r(_max_heap);
            std::sort(res_l.begin(), res_l.end(), _comp);
            std::sort(res_r.begin(), res_r.end(), _comp);
            res_l.reserve(size());
            std::copy(res_r.begin(), res_r.end(), std::back_inserter(res_l));
            return res_l;
        }

    private:
        compare_type _comp;
        struct {
            compare_type* comp;
            bool operator()(const value_type& x, const value_type& y) { return (*comp)(y, x); }
        } _rev_comp{ &_comp };

        std::vector<value_type> _max_heap, _min_heap;
        value_type pivot;

        void ensure_min_heap_nonempty() {
            const int siz = size();
            assert(siz);
            if (not _min_heap.empty()) return;
            if (siz == 1) {
                std::swap(_min_heap, _max_heap);
                pivot = _min_heap.front();
            } else {
                const int mid = (siz + 1) >> 1;
                std::nth_element(_max_heap.begin(), _max_heap.begin() + mid - 1, _max_heap.end(), _comp);
                pivot = _max_heap[mid - 1];
                _min_heap.reserve(mid);
                std::move(_max_heap.begin(), _max_heap.begin() + mid, std::back_inserter(_min_heap));
                _max_heap.erase(_max_heap.begin(), _max_heap.begin() + mid);
                std::make_heap(_max_heap.begin(), _max_heap.end(), _comp);
                std::make_heap(_min_heap.begin(), _min_heap.end(), _rev_comp);
            }
        }
        void ensure_max_heap_nonempty() {
            const int siz = size();
            assert(siz);
            if (not _max_heap.empty()) return;
            if (siz == 1) {
                std::swap(_min_heap, _max_heap);
            } else {
                const int mid = siz >> 1;
                std::nth_element(_min_heap.begin(), _min_heap.begin() + mid - 1, _min_heap.end(), _comp);
                pivot = _min_heap[mid - 1];
                _max_heap.reserve(siz - mid);
                std::move(_min_heap.begin() + mid, _min_heap.end(), std::back_inserter(_max_heap));
                _min_heap.erase(_min_heap.begin() + mid, _min_heap.end());
                std::make_heap(_max_heap.begin(), _max_heap.end(), _comp);
                std::make_heap(_min_heap.begin(), _min_heap.end(), _rev_comp);
            }
        }
    };
} // namespace suisen



#line 6 "test/src/datastructure/double_ended_priority_queue/double_ended_priority_queue.test.cpp"

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

    int n, q;
    std::cin >> n >> q;

    std::vector<int> init(n);
    for (auto &e : init) std::cin >> e;
    suisen::DoubleEndedPriorityQueue<int> pq(init.begin(), init.end());

    for (int i = 0; i < q; ++i) {
        int query_type;
        std::cin >> query_type;
        if (query_type == 0) {
            int x;
            std::cin >> x;
            pq.push(x);
        } else if (query_type == 1) {
            std::cout << pq.pop_min() << '\n';
        } else {
            std::cout << pq.pop_max() << '\n';
        }
    }
    return 0;
}

Back to top page