shogo-library

競技プログラミング用の C++ ライブラリ

View the Project on GitHub shogo314/shogo-library

:heavy_check_mark: 上位 K 個の値とその個数
(src/data_structures/priority_counter.hpp)

値を追加していき、大きい方から K 種類の値とそれぞれの個数だけを保持します。+ で2つを合わせられるので、セグメント木に乗せて「区間の最大値とその個数」「区間で2番目に大きい値」などを求めるのに使えます。

MaxCounter<int, 2> pc;                 // 大きい方から 2 種類
for (int x : {3, 5, 1, 5, 4}) pc.inc(x);
// pc.data = { {5, 2}, {4, 1} }

MaxCounter<int> a(3), b(3);
auto c = a + b;                        // c.data[0] = {3, 2}

型

型 説明
PriorityCounter<T, K = 1, Compare = std::less<T>> Compare で大きい方から K 種類を保持
MaxCounter<T, K = 1> 大きい方から K 種類
MinCounter<T, K = 1> 小さい方から K 種類

メンバ

メンバ 説明 計算量
PriorityCounter() 空 $O(K)$
PriorityCounter(x) 値 x が 1 個 $O(K)$
inc(x, c = 1) 値 x を c 個追加する。上位 K 種類に入らなければ捨てる $O(K)$
a + b、a += b 合わせる $O(K^2)$
size() 保持している値の種類数 (K 以下) $O(K)$
data[i] 上位から i 番目の {値, 個数}。i >= size() では個数が 0  

値は追加されるだけなので、上位 K 種類から外れた値が再び上位に入ることはありません。そのため、追加の順番によらず上位 K 種類とその個数は正しく求まります。c は正の値にしてください (個数を減らす使い方はできません)。

Verified with

Code

#pragma once

/**
 * @file priority_counter.hpp
 * @brief 上位 K 個の値とその個数
 */

#include <array>
#include <cstddef>
#include <functional>
#include <map>
#include <type_traits>
#include <utility>

template <typename T, size_t K = 1, class Compare = std::less<T>, std::enable_if_t<(K >= 1), std::nullptr_t> = nullptr>
struct PriorityCounter {
    using key_type = T;
    using mapped_type = long long;
    using value_type = std::pair<key_type, mapped_type>;
    inline constexpr static size_t k = K;

    // C++17 の constexpr コンストラクタではすべてのメンバの初期化が必要なので、値初期化しておく
    std::array<value_type, K> data{};

    constexpr PriorityCounter() {}
    constexpr PriorityCounter(const key_type& x) : PriorityCounter() {
        data[0].first = x;
        data[0].second = 1;
    }

    constexpr size_t size() const {
        for (size_t i = 0; i < k; i++) {
            if (data[i].second == 0) {
                return i;
            }
        }
        return k;
    }

    constexpr void inc(const key_type& x, const mapped_type& c = 1) {
        bool f = false;
        key_type pre_key{};
        mapped_type pre_mapped{};
        for (size_t i = 0; i < k; i++) {
            if (f) {
                // std::swap は C++17 では constexpr でないので手で入れ替える
                key_type tmp_key = data[i].first;
                data[i].first = pre_key;
                pre_key = tmp_key;
                mapped_type tmp_mapped = data[i].second;
                data[i].second = pre_mapped;
                pre_mapped = tmp_mapped;
                continue;
            }
            if (data[i].second == 0) {
                data[i].first = x;
                data[i].second += c;
                return;
            }
            if (Compare()(data[i].first, x)) {
                pre_key = data[i].first;
                pre_mapped = data[i].second;
                data[i].first = x;
                data[i].second = c;
                f = true;
            } else if (not Compare()(x, data[i].first)) {
                data[i].second += c;
                return;
            }
        }
    }

    constexpr PriorityCounter& operator+=(const PriorityCounter& o) {
        for (size_t i = 0; i < k; i++) {
            if (o.data[i].second == 0) break;
            this->inc(o.data[i].first, o.data[i].second);
        }
        return *this;
    }
    constexpr friend PriorityCounter operator+(const PriorityCounter& a, const PriorityCounter& b) {
        PriorityCounter res(a);
        return res += b;
    }
};

template <typename T, size_t K = 1>
using MaxCounter = PriorityCounter<T, K>;

template <typename T, size_t K = 1>
using MinCounter = PriorityCounter<T, K, std::greater<T>>;
#line 2 "src/data_structures/priority_counter.hpp"

/**
 * @file priority_counter.hpp
 * @brief 上位 K 個の値とその個数
 */

#include <array>
#include <cstddef>
#include <functional>
#include <map>
#include <type_traits>
#include <utility>

template <typename T, size_t K = 1, class Compare = std::less<T>, std::enable_if_t<(K >= 1), std::nullptr_t> = nullptr>
struct PriorityCounter {
    using key_type = T;
    using mapped_type = long long;
    using value_type = std::pair<key_type, mapped_type>;
    inline constexpr static size_t k = K;

    // C++17 の constexpr コンストラクタではすべてのメンバの初期化が必要なので、値初期化しておく
    std::array<value_type, K> data{};

    constexpr PriorityCounter() {}
    constexpr PriorityCounter(const key_type& x) : PriorityCounter() {
        data[0].first = x;
        data[0].second = 1;
    }

    constexpr size_t size() const {
        for (size_t i = 0; i < k; i++) {
            if (data[i].second == 0) {
                return i;
            }
        }
        return k;
    }

    constexpr void inc(const key_type& x, const mapped_type& c = 1) {
        bool f = false;
        key_type pre_key{};
        mapped_type pre_mapped{};
        for (size_t i = 0; i < k; i++) {
            if (f) {
                // std::swap は C++17 では constexpr でないので手で入れ替える
                key_type tmp_key = data[i].first;
                data[i].first = pre_key;
                pre_key = tmp_key;
                mapped_type tmp_mapped = data[i].second;
                data[i].second = pre_mapped;
                pre_mapped = tmp_mapped;
                continue;
            }
            if (data[i].second == 0) {
                data[i].first = x;
                data[i].second += c;
                return;
            }
            if (Compare()(data[i].first, x)) {
                pre_key = data[i].first;
                pre_mapped = data[i].second;
                data[i].first = x;
                data[i].second = c;
                f = true;
            } else if (not Compare()(x, data[i].first)) {
                data[i].second += c;
                return;
            }
        }
    }

    constexpr PriorityCounter& operator+=(const PriorityCounter& o) {
        for (size_t i = 0; i < k; i++) {
            if (o.data[i].second == 0) break;
            this->inc(o.data[i].first, o.data[i].second);
        }
        return *this;
    }
    constexpr friend PriorityCounter operator+(const PriorityCounter& a, const PriorityCounter& b) {
        PriorityCounter res(a);
        return res += b;
    }
};

template <typename T, size_t K = 1>
using MaxCounter = PriorityCounter<T, K>;

template <typename T, size_t K = 1>
using MinCounter = PriorityCounter<T, K, std::greater<T>>;
Back to top page