top_k_sum.hpp¶

SECTIONData Structure INCLUDEnoya/top_k_sum.hpp

动态维护集合中最大的或最小的 \(k\) 个元素之和；适合排名阈值变化、选取固定数量最优元素的题目。

Complexity: Time: O(log n) insert/erase/set-k and O(1) sum query. Space: O(n).

Implementation¶

ReferenceSourceExpanded

当前头文件，省略 include guard；依赖见 #include。

/// @complexity Time: O(log n) insert/erase/set-k and O(1) sum query.
/// Space: O(n).

#include <algorithm>
#include <cassert>
#include <set>
#include <utility>

namespace noya {

/// @brief Dynamic multiset maintaining the sum of the k largest values in
/// O(log n) per insert, erase, or k change.
template <class T> struct top_k_sum {
  using entry = std::pair<T, int>;
  int k = 0;
  int tot = 0;
  T sum{};
  std::set<entry> hi;
  std::set<entry> lo;

  top_k_sum() = default;
  explicit top_k_sum(int k_) : k(k_) { assert(k >= 0); }

  /// @brief Insert value and return a stable handle for exact erasure.
  int insert(const T &val) {
    int id = tot++;
    lo.emplace(val, id);
    rebalance();
    return id;
  }

  /// @brief Erase the entry identified by (value, handle).
  void erase(const T &val, int id) {
    entry tar{val, id};
    if (auto it = hi.find(tar); it != hi.end()) {
      sum -= it->first;
      hi.erase(it);
    } else {
      auto ri = lo.find(tar);
      assert(ri != lo.end());
      lo.erase(ri);
    }
    rebalance();
  }

  void set_k(int k_) {
    assert(k_ >= 0);
    k = k_;
    rebalance();
  }

  const T &query() const { return sum; }
  int size() const { return int(hi.size() + lo.size()); }

private:
  void rebalance() {
    int tar = std::min(k, size());
    while (int(hi.size()) < tar) {
      auto it = std::prev(lo.end());
      sum += it->first;
      hi.insert(*it);
      lo.erase(it);
    }
    while (int(hi.size()) > tar) {
      auto it = hi.begin();
      sum -= it->first;
      lo.insert(*it);
      hi.erase(it);
    }
    while (!hi.empty() && !lo.empty() &&
           *hi.begin() < *lo.rbegin()) {
      entry low = *hi.begin();
      entry big = *lo.rbegin();
      hi.erase(hi.begin());
      lo.erase(std::prev(lo.end()));
      hi.insert(big);
      lo.insert(low);
      sum += big.first - low.first;
    }
  }
};

} // namespace noya

#ifndef NOYA_TOP_K_SUM_HPP
#define NOYA_TOP_K_SUM_HPP 1

/// @complexity Time: O(log n) insert/erase/set-k and O(1) sum query.
/// Space: O(n).

#include <algorithm>
#include <cassert>
#include <set>
#include <utility>

namespace noya {

/// @brief Dynamic multiset maintaining the sum of the k largest values in
/// O(log n) per insert, erase, or k change.
template <class T> struct top_k_sum {
  using entry = std::pair<T, int>;
  int k = 0;
  int tot = 0;
  T sum{};
  std::set<entry> hi;
  std::set<entry> lo;

  top_k_sum() = default;
  explicit top_k_sum(int k_) : k(k_) { assert(k >= 0); }

  /// @brief Insert value and return a stable handle for exact erasure.
  int insert(const T &val) {
    int id = tot++;
    lo.emplace(val, id);
    rebalance();
    return id;
  }

  /// @brief Erase the entry identified by (value, handle).
  void erase(const T &val, int id) {
    entry tar{val, id};
    if (auto it = hi.find(tar); it != hi.end()) {
      sum -= it->first;
      hi.erase(it);
    } else {
      auto ri = lo.find(tar);
      assert(ri != lo.end());
      lo.erase(ri);
    }
    rebalance();
  }

  void set_k(int k_) {
    assert(k_ >= 0);
    k = k_;
    rebalance();
  }

  const T &query() const { return sum; }
  int size() const { return int(hi.size() + lo.size()); }

private:
  void rebalance() {
    int tar = std::min(k, size());
    while (int(hi.size()) < tar) {
      auto it = std::prev(lo.end());
      sum += it->first;
      hi.insert(*it);
      lo.erase(it);
    }
    while (int(hi.size()) > tar) {
      auto it = hi.begin();
      sum -= it->first;
      lo.insert(*it);
      hi.erase(it);
    }
    while (!hi.empty() && !lo.empty() &&
           *hi.begin() < *lo.rbegin()) {
      entry low = *hi.begin();
      entry big = *lo.rbegin();
      hi.erase(hi.begin());
      lo.erase(std::prev(lo.end()));
      hi.insert(big);
      lo.insert(low);
      sum += big.first - low.first;
    }
  }
};

} // namespace noya

#endif // NOYA_TOP_K_SUM_HPP

#include <algorithm>
#include <cassert>
#include <set>
#include <utility>

/// @complexity Time: O(log n) insert/erase/set-k and O(1) sum query.
/// Space: O(n).

namespace noya {

/// @brief Dynamic multiset maintaining the sum of the k largest values in
/// O(log n) per insert, erase, or k change.
template <class T> struct top_k_sum {
  using entry = std::pair<T, int>;
  int k = 0;
  int tot = 0;
  T sum{};
  std::set<entry> hi;
  std::set<entry> lo;

  top_k_sum() = default;
  explicit top_k_sum(int k_) : k(k_) { assert(k >= 0); }

  /// @brief Insert value and return a stable handle for exact erasure.
  int insert(const T &val) {
    int id = tot++;
    lo.emplace(val, id);
    rebalance();
    return id;
  }

  /// @brief Erase the entry identified by (value, handle).
  void erase(const T &val, int id) {
    entry tar{val, id};
    if (auto it = hi.find(tar); it != hi.end()) {
      sum -= it->first;
      hi.erase(it);
    } else {
      auto ri = lo.find(tar);
      assert(ri != lo.end());
      lo.erase(ri);
    }
    rebalance();
  }

  void set_k(int k_) {
    assert(k_ >= 0);
    k = k_;
    rebalance();
  }

  const T &query() const { return sum; }
  int size() const { return int(hi.size() + lo.size()); }

private:
  void rebalance() {
    int tar = std::min(k, size());
    while (int(hi.size()) < tar) {
      auto it = std::prev(lo.end());
      sum += it->first;
      hi.insert(*it);
      lo.erase(it);
    }
    while (int(hi.size()) > tar) {
      auto it = hi.begin();
      sum -= it->first;
      lo.insert(*it);
      hi.erase(it);
    }
    while (!hi.empty() && !lo.empty() &&
           *hi.begin() < *lo.rbegin()) {
      entry low = *hi.begin();
      entry big = *lo.rbegin();
      hi.erase(hi.begin());
      lo.erase(std::prev(lo.end()));
      hi.insert(big);
      lo.insert(low);
      sum += big.first - low.first;
    }
  }
};

} // namespace noya