昇順 $K$ 個の和を取れるpriority_queue
(Src/DataStructure/Other/PriorityProductSet.hpp)

• View this file on GitHub
• Last update: 2025-04-17 19:44:48+09:00
• Include: #include "Src/DataStructure/Other/PriorityProductSet.hpp"

概要

要素の挿入と昇順 $K$ 要素の和が取れるデータ構造。要素の削除は $K$ 番目からしか行えないことに注意。

• 群を載せることができる

ライブラリ

PriorityProductSet()
PriorityProductSet(usize K)
void insert(V&& v)
void insert(const V& v)
usize size() const
inline usize K() const
std::optional<V> product() const
std::optional<V> productRemain() const
V productAll() const
V popK()

Depends on

• Src/Algebra/Group/GroupConcept.hpp
• Src/Algebra/Monoid/MonoidConcept.hpp
• Src/Algebra/Semigroup/SemigroupConcept.hpp
• 標準データ型のエイリアス (Src/Template/TypeAlias.hpp)

Verified with

• Test/AtCoder/arc196_a.test.cpp
• Test/yukicoder/649.test.cpp

Code

#pragma once

#include "../../Template/TypeAlias.hpp"
#include "../../Algebra/Group/GroupConcept.hpp"

#include <cassert>
#include <optional>
#include <utility>
#include <queue>

namespace zawa {

// 1-indexed
template <concepts::Group G>
class PriorityProductSet {
public:

    using V = G::Element;

    PriorityProductSet() = default;

    PriorityProductSet(usize K) : m_K{K} {}

    void insert(V&& v) {
        pushSmall(std::move(v));
        adjust();
    }

    void insert(const V& v) {
        insert(V{v});
    }

    usize size() const {
        return m_small.size() + m_big.size();
    }

    inline usize K() const {
        return m_K;
    }

    void setK(usize K) {
        m_K = K;
        adjust();
    }

    std::optional<V> product() const {
        return m_small.size() == m_K ? std::optional<V>{m_sv} : std::nullopt;
    }

    std::optional<V> productRemain() const {
        return m_small.size() == m_K ? std::optional<V>{m_bv} : std::nullopt;
    }

    V productAll() const {
        return G::operation(m_sv, m_bv);
    }

    V popK() {
        assert(m_K >= 1u and size() >= m_K);
        V res = popSmall(); 
        adjust();
        return res;
    }

private:

    std::priority_queue<V> m_small;

    std::priority_queue<V, std::vector<V>, std::greater<V>> m_big;

    V m_sv = G::identity(), m_bv = G::identity();

    usize m_K = 0;

    void pushSmall(V&& v) {
        m_sv = G::operation(m_sv, v);
        m_small.push(std::move(v));
    }

    V popSmall() {
        assert(m_small.size());
        V res = m_small.top();
        m_small.pop();
        m_sv = G::operation(m_sv, G::inverse(res));
        return res;
    }

    void pushBig(V&& v) {
        m_bv = G::operation(m_bv, v);
        m_big.push(std::move(v));
    }

    V popBig() {
        assert(m_big.size());
        V res = m_big.top();
        m_big.pop();
        m_bv = G::operation(m_bv, G::inverse(res));
        return res;
    }

    void adjust() {
        while (m_small.size() > m_K) pushBig(popSmall());
        while (m_small.size() < m_K and m_big.size()) pushSmall(popBig());
    }
};

} // namespace zawa

#line 2 "Src/DataStructure/Other/PriorityProductSet.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 2 "Src/Algebra/Group/GroupConcept.hpp"

#line 2 "Src/Algebra/Monoid/MonoidConcept.hpp"

#line 2 "Src/Algebra/Semigroup/SemigroupConcept.hpp"

#include <concepts>

namespace zawa {

namespace concepts {

template <class T>
concept Semigroup = requires {
    typename T::Element;
    { T::operation(std::declval<typename T::Element>(), std::declval<typename T::Element>()) } -> std::same_as<typename T::Element>;
};

} // namespace concepts

} // namespace zawa
#line 4 "Src/Algebra/Monoid/MonoidConcept.hpp"

#line 6 "Src/Algebra/Monoid/MonoidConcept.hpp"

namespace zawa {

namespace concepts {

template <class T>
concept Identitiable = requires {
    typename T::Element;
    { T::identity() } -> std::same_as<typename T::Element>;
};

template <class T>
concept Monoid = Semigroup<T> and Identitiable<T>;

} // namespace

} // namespace zawa
#line 4 "Src/Algebra/Group/GroupConcept.hpp"

namespace zawa {

namespace concepts {

template <class T>
concept Inversible = requires {
    typename T::Element;
    { T::inverse(std::declval<typename T::Element>()) } -> std::same_as<typename T::Element>;
};

template <class T>
concept Group = Monoid<T> and Inversible<T>;

} // namespace Concept

} // namespace zawa
#line 5 "Src/DataStructure/Other/PriorityProductSet.hpp"

#include <cassert>
#include <optional>
#include <utility>
#include <queue>

namespace zawa {

// 1-indexed
template <concepts::Group G>
class PriorityProductSet {
public:

    using V = G::Element;

    PriorityProductSet() = default;

    PriorityProductSet(usize K) : m_K{K} {}

    void insert(V&& v) {
        pushSmall(std::move(v));
        adjust();
    }

    void insert(const V& v) {
        insert(V{v});
    }

    usize size() const {
        return m_small.size() + m_big.size();
    }

    inline usize K() const {
        return m_K;
    }

    void setK(usize K) {
        m_K = K;
        adjust();
    }

    std::optional<V> product() const {
        return m_small.size() == m_K ? std::optional<V>{m_sv} : std::nullopt;
    }

    std::optional<V> productRemain() const {
        return m_small.size() == m_K ? std::optional<V>{m_bv} : std::nullopt;
    }

    V productAll() const {
        return G::operation(m_sv, m_bv);
    }

    V popK() {
        assert(m_K >= 1u and size() >= m_K);
        V res = popSmall(); 
        adjust();
        return res;
    }

private:

    std::priority_queue<V> m_small;

    std::priority_queue<V, std::vector<V>, std::greater<V>> m_big;

    V m_sv = G::identity(), m_bv = G::identity();

    usize m_K = 0;

    void pushSmall(V&& v) {
        m_sv = G::operation(m_sv, v);
        m_small.push(std::move(v));
    }

    V popSmall() {
        assert(m_small.size());
        V res = m_small.top();
        m_small.pop();
        m_sv = G::operation(m_sv, G::inverse(res));
        return res;
    }

    void pushBig(V&& v) {
        m_bv = G::operation(m_bv, v);
        m_big.push(std::move(v));
    }

    V popBig() {
        assert(m_big.size());
        V res = m_big.top();
        m_big.pop();
        m_bv = G::operation(m_bv, G::inverse(res));
        return res;
    }

    void adjust() {
        while (m_small.size() > m_K) pushBig(popSmall());
        while (m_small.size() < m_K and m_big.size()) pushSmall(popBig());
    }
};

} // namespace zawa

Back to top page