Priority Sum Structure (structure/others/priority-sum-structure.hpp)

• View this file on GitHub
• View document part on GitHub
• Last update: 2026-08-04 11:55:02+00:00
• Include: #include "structure/others/priority-sum-structure.hpp"

PrioritySumStructure は、要素の追加・削除を行いながら、優先度の高い $k$ 個の総和を管理するデータ構造です。

MaximumSum は大きい方から $k$ 個、MinimumSum は小さい方から $k$ 個の総和を管理します。

template <typename T>
using MaximumSum = PrioritySumStructure<T, std::greater<T>, std::less<T>>;

template <typename T>
using MinimumSum = PrioritySumStructure<T, std::less<T>, std::greater<T>>;

コンストラクタ

PrioritySumStructure<T>(int k)

優先度の高い $k$ 個の総和を管理するデータ構造を初期化します。

計算量

• $O(1)$

query

T query() const

現在の要素のうち、優先度の高い $k$ 個の総和を返します。

MaximumSum では大きい方から $k$ 個、MinimumSum では小さい方から $k$ 個の総和を返します。

計算量

• $O(1)$

kth_element

T kth_element()

現在の $k$ 番目 (1-indexed) の要素を返します。

MaximumSum では $k$ 番目に大きい値、MinimumSum では $k$ 番目に小さい値を返します。

計算量

• amortized $O(\log n)$

insert

void insert(T x)

要素 $x$ を追加します。

計算量

• amortized $O(\log n)$

erase

void erase(T x)

要素 $x$ を 1 つ削除します。

制約

• $x$ が存在する

計算量

• amortized $O(\log n)$

set_k

void set_k(std::size_t kk)

管理する要素の個数を kk に変更します。

計算量

• amortized $O(\log n)$

get_k

std::size_t get_k() const

現在の $k$ を返します。

計算量

• $O(1)$

size

std::size_t size() const

現在の要素数を返します。

計算量

• $O(1)$

Verified with

• test/unittest/priority-sum-structure.test.cpp

Code

#pragma once

#include <cassert>
#include <cstddef>
#include <functional>
#include <queue>
#include <vector>

template <typename T, typename Compare = std::less<T>,
          typename RCompare = std::greater<T>>
struct PrioritySumStructure {
  std::size_t k;
  T sum;

  std::priority_queue<T, std::vector<T>, Compare> in, d_in;
  std::priority_queue<T, std::vector<T>, RCompare> out, d_out;

  PrioritySumStructure(int k) : k(k), sum(0) {}

  void modify() {
    while (in.size() - d_in.size() < k && !out.empty()) {
      auto p = out.top();
      out.pop();
      if (!d_out.empty() && p == d_out.top()) {
        d_out.pop();
      } else {
        sum += p;
        in.emplace(p);
      }
    }
    while (in.size() - d_in.size() > k) {
      auto p = in.top();
      in.pop();
      if (!d_in.empty() && p == d_in.top()) {
        d_in.pop();
      } else {
        sum -= p;
        out.emplace(p);
      }
    }
    while (!d_in.empty() && in.top() == d_in.top()) {
      in.pop();
      d_in.pop();
    }
  }

  T query() const { return sum; }

  T kth_element() {
    assert(0 < k && k <= size());
    modify();
    return in.top();
  }

  void insert(T x) {
    in.emplace(x);
    sum += x;
    modify();
  }

  void erase(T x) {
    assert(size());
    if (!in.empty() && in.top() == x) {
      sum -= x;
      in.pop();
    } else if (!in.empty() && RCompare()(in.top(), x)) {
      sum -= x;
      d_in.emplace(x);
    } else {
      d_out.emplace(x);
    }
    modify();
  }

  void set_k(std::size_t kk) {
    k = kk;
    modify();
  }

  std::size_t get_k() const { return k; }

  std::size_t size() const {
    return in.size() + out.size() - d_in.size() - d_out.size();
  }
};

template <typename T>
using MaximumSum = PrioritySumStructure<T, std::greater<T>, std::less<T>>;

template <typename T>
using MinimumSum = PrioritySumStructure<T, std::less<T>, std::greater<T>>;

#line 2 "structure/others/priority-sum-structure.hpp"

#include <cassert>
#include <cstddef>
#include <functional>
#include <queue>
#include <vector>

template <typename T, typename Compare = std::less<T>,
          typename RCompare = std::greater<T>>
struct PrioritySumStructure {
  std::size_t k;
  T sum;

  std::priority_queue<T, std::vector<T>, Compare> in, d_in;
  std::priority_queue<T, std::vector<T>, RCompare> out, d_out;

  PrioritySumStructure(int k) : k(k), sum(0) {}

  void modify() {
    while (in.size() - d_in.size() < k && !out.empty()) {
      auto p = out.top();
      out.pop();
      if (!d_out.empty() && p == d_out.top()) {
        d_out.pop();
      } else {
        sum += p;
        in.emplace(p);
      }
    }
    while (in.size() - d_in.size() > k) {
      auto p = in.top();
      in.pop();
      if (!d_in.empty() && p == d_in.top()) {
        d_in.pop();
      } else {
        sum -= p;
        out.emplace(p);
      }
    }
    while (!d_in.empty() && in.top() == d_in.top()) {
      in.pop();
      d_in.pop();
    }
  }

  T query() const { return sum; }

  T kth_element() {
    assert(0 < k && k <= size());
    modify();
    return in.top();
  }

  void insert(T x) {
    in.emplace(x);
    sum += x;
    modify();
  }

  void erase(T x) {
    assert(size());
    if (!in.empty() && in.top() == x) {
      sum -= x;
      in.pop();
    } else if (!in.empty() && RCompare()(in.top(), x)) {
      sum -= x;
      d_in.emplace(x);
    } else {
      d_out.emplace(x);
    }
    modify();
  }

  void set_k(std::size_t kk) {
    k = kk;
    modify();
  }

  std::size_t get_k() const { return k; }

  std::size_t size() const {
    return in.size() + out.size() - d_in.size() - d_out.size();
  }
};

template <typename T>
using MaximumSum = PrioritySumStructure<T, std::greater<T>, std::less<T>>;

template <typename T>
using MinimumSum = PrioritySumStructure<T, std::less<T>, std::greater<T>>;

Back to top page