Src/DataStructure/Heap/EraseablePriorityQueue.hpp

View this file on GitHub
Last update: 2026-01-14 16:33:38+09:00
Include: #include "Src/DataStructure/Heap/EraseablePriorityQueue.hpp"

Depends on

Required by

要素昇順 $K$ 個、降順 $N - K$ 個の総積を管理 (Src/DataStructure/Heap/PartitionedProducts.hpp)

Verified with

Code

#pragma once

#include "BinaryHeap.hpp"

namespace zawa {

template <class T, class Comp = std::less<T>>
requires std::strict_weak_order<Comp, const T&, const T&>
class EraseablePriorityQueue {
private:

    BinaryHeap<T, Comp> m_que, m_rm;

    usize m_cnt = 0;

public:

    inline usize size() const {
        return m_cnt;
    }

    inline bool empty() const {
        return m_cnt == 0;
    }

    EraseablePriorityQueue(Comp comp = {}) 
        : m_que{comp}, m_rm{comp}, m_cnt{0} {}

    template <std::forward_iterator It>
    requires std::same_as<std::iter_value_t<It>, T>
    EraseablePriorityQueue(It first, It last, Comp comp = {})
        : m_que{first, last, comp}, m_rm{comp}, m_cnt{m_que.size()} {}

    EraseablePriorityQueue(std::vector<T> a, Comp comp = {}) 
        : m_que{a, comp}, m_rm{comp}, m_cnt{m_que.size()} {}

    template <class U>
    requires std::same_as<std::remove_cvref_t<U>, T>
    void push(U&& v) {
        m_que.push(std::forward<U>(v));
        m_cnt++;
    }

    template <class U>
    requires std::same_as<std::remove_cvref_t<U>, T>
    void erase(U&& v) {
        assert(size() and "HeapUnderFlow");
        m_rm.push(std::forward<U>(v));
        m_cnt--;
    }

    const T& top() {
        assert(size() and "HeapUnderFlow");
        normalize();
        return m_que.top();
    }

    T pop() {
        assert(size() and "HeapUnderFlow");
        normalize();
        T res = m_que.top();
        m_que.pop();
        m_cnt--;
        return res;
    }

    std::vector<T> container() const {
        BinaryHeap que = m_que, rm = m_rm;  
        std::vector<T> res;
        while (que.size()) {
            if (rm.empty() or que.comp()(que.top(), rm.top())) {
                res.push_back(que.top());
                que.pop();
            }
            else if (que.top() == rm.top())
                que.pop(), rm.pop();
            else
                rm.pop();
        }
        return res;
    }

private:

    void normalize() {
        while (m_rm.size() and m_que.size()) {
            if (m_que.top() == m_rm.top())
                m_que.pop(), m_rm.pop();
            else if (m_que.comp()(m_rm.top(), m_que.top()))
                m_rm.pop();
            else
                break;
        }
    }
};

} // namespace zawa

#line 2 "Src/DataStructure/Heap/EraseablePriorityQueue.hpp"

#line 2 "Src/DataStructure/Heap/BinaryHeap.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/DataStructure/Heap/BinaryHeap.hpp"

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

namespace zawa {

template <class T, class Comp = std::less<T>>
requires std::strict_weak_order<Comp, const T&, const T&>
class BinaryHeap {
private:

    Comp m_comp;

    std::vector<T> m_dat;

public:

    inline usize size() const {
        return m_dat.size() - 1;
    }

    inline bool empty() const {
        return m_dat.size() == 1;
    }

    inline const Comp& comp() const {
        return m_comp;
    }

    using const_iterator = typename decltype(m_dat)::const_iterator;

    const_iterator begin() const {
        return m_dat.begin() + 1;
    }

    const_iterator end() const {
        return m_dat.end();
    }

    BinaryHeap(Comp comp = {}) 
        : m_comp{comp}, m_dat(1) {}

    template <std::forward_iterator It>
    requires std::same_as<std::iter_value_t<It>, T>
    BinaryHeap(It first, It last, Comp comp = {}) 
        : m_comp{comp}, m_dat(1) {
        m_dat.insert(m_dat.end(), first, last);
        build();
    }

    BinaryHeap(std::vector<T>&& a, Comp comp = {}) 
        : m_comp{comp}, m_dat(a.size() + 1) {
        std::ranges::copy(std::make_move_iterator(a.begin()), std::make_move_iterator(a.end()), m_dat.begin() + 1);
        build();
    }

    BinaryHeap(const std::vector<T>& a, Comp comp = {}) 
        : m_comp{comp}, m_dat(a.size() + 1) {
        std::ranges::copy(a.begin(), a.end(), m_dat.begin() + 1);
        build();
    }

    const T& top() const {
        assert(size() and "HeapUnderFlow");
        return m_dat[1];
    }

    void push(T&& v) {
        m_dat.push_back(std::move(v));
        upHeap(size());
    }

    void push(const T& v) {
        m_dat.push_back(v);
        upHeap(size());
    }

    void pop() {
        assert(size() and "HeapUnderFlow");
        if (size() > 1)
            std::swap(m_dat[1], m_dat.back());
        m_dat.pop_back();
        if (size() > 1)
            downHeap(1, size());
    }

private:

    void build() {
        const usize n = size();
        for (usize i = (n >> 1) ; i ; i--) 
            downHeap(i, n);
    }

    void upHeap(usize i) {
        while (i >> 1 and m_comp(m_dat[i], m_dat[i >> 1])) {
            std::swap(m_dat[i], m_dat[i >> 1]);
            i >>= 1;
        }
    }

    void downHeap(usize i, usize n) {
        while ((i << 1) <= n) {
            usize j = i << 1;
            if (j + 1 <= n and m_comp(m_dat[j + 1], m_dat[j]))
                j++;
            if (!m_comp(m_dat[j], m_dat[i]))
                break;
            std::swap(m_dat[i], m_dat[j]);
            i = j;
        }
    }
};

} // namespace zawa
#line 4 "Src/DataStructure/Heap/EraseablePriorityQueue.hpp"

namespace zawa {

template <class T, class Comp = std::less<T>>
requires std::strict_weak_order<Comp, const T&, const T&>
class EraseablePriorityQueue {
private:

    BinaryHeap<T, Comp> m_que, m_rm;

    usize m_cnt = 0;

public:

    inline usize size() const {
        return m_cnt;
    }

    inline bool empty() const {
        return m_cnt == 0;
    }

    EraseablePriorityQueue(Comp comp = {}) 
        : m_que{comp}, m_rm{comp}, m_cnt{0} {}

    template <std::forward_iterator It>
    requires std::same_as<std::iter_value_t<It>, T>
    EraseablePriorityQueue(It first, It last, Comp comp = {})
        : m_que{first, last, comp}, m_rm{comp}, m_cnt{m_que.size()} {}

    EraseablePriorityQueue(std::vector<T> a, Comp comp = {}) 
        : m_que{a, comp}, m_rm{comp}, m_cnt{m_que.size()} {}

    template <class U>
    requires std::same_as<std::remove_cvref_t<U>, T>
    void push(U&& v) {
        m_que.push(std::forward<U>(v));
        m_cnt++;
    }

    template <class U>
    requires std::same_as<std::remove_cvref_t<U>, T>
    void erase(U&& v) {
        assert(size() and "HeapUnderFlow");
        m_rm.push(std::forward<U>(v));
        m_cnt--;
    }

    const T& top() {
        assert(size() and "HeapUnderFlow");
        normalize();
        return m_que.top();
    }

    T pop() {
        assert(size() and "HeapUnderFlow");
        normalize();
        T res = m_que.top();
        m_que.pop();
        m_cnt--;
        return res;
    }

    std::vector<T> container() const {
        BinaryHeap que = m_que, rm = m_rm;  
        std::vector<T> res;
        while (que.size()) {
            if (rm.empty() or que.comp()(que.top(), rm.top())) {
                res.push_back(que.top());
                que.pop();
            }
            else if (que.top() == rm.top())
                que.pop(), rm.pop();
            else
                rm.pop();
        }
        return res;
    }

private:

    void normalize() {
        while (m_rm.size() and m_que.size()) {
            if (m_que.top() == m_rm.top())
                m_que.pop(), m_rm.pop();
            else if (m_que.comp()(m_rm.top(), m_que.top()))
                m_rm.pop();
            else
                break;
        }
    }
};

} // namespace zawa