如何使用std::priority_queue创建固定大小的最小堆？

Question

我可以将最小堆定义为：

priority_queue<int, vector<int>, greater> pq;

我有一个整数流。最小堆的大小是一个固定值k。似乎priority_queue不能做到这一点。

Answer 1

如果您想使用std::priority_queue，那么将队列的大小限制为k元素是微不足道的。但是请注意，您需要使用最大堆，而不是最小堆，因为您需要知道是否应该将新到达的值插入到堆中，如果新到达的值小于堆中当前的最大值，则会发生这种情况。

class Topk {
  public:
    Topk(int k) : k_(k) {}
    void insert(int value) {
      if (q_.size() < k_) q_.push(value);
      else if (value < q_.top()) { q_.pop(); q_.push(value); }
    }
    std::vector<int> finalize() {
      std::vector<int> result(q_.size());
      while (q_.size()) {
        result[q_.size() - 1] = q_.top();
        q_.pop();
      }
      return result;
    }
  private:
    int k_;
    std::priority_queue<int> q_;
}

仅仅使用堆算法实际上并不复杂：

class Topk {
  public:
    Topk(int k) : k_(k) {}
    void insert(int value) {
      if (c_.size() < k_) {
        c_.push_back(value);
        if (c_.size() == k_) make_heap(c_.begin(), c_.end());
      }
      else if (value < c_[0]) {
        /* See note below */
        pop_heap(c_.begin(), c_.end());
        c_.back() = value;
        push_heap(c_.begin(), c_.end());
      }
    }
    std::vector<int> finalize() {
      if (c_.size() < k_)
        std::sort(c_.begin(), c_.end());
      else
        sort_heap(c_.begin(), c_end());
      std::vector<int> c;
      std::swap(c, c_);
      return std::move(c);
    }
  private:
    /* invariant: if c_.size() == k, then c_ is a maxheap. */
    int k_;
    std::vector<int> c_;
}

注意：<algorithm>不包括heap_sift_down操作，这对此应用程序来说是不幸的；弹出/交换/推送操作可以替换为交换/ sift_down操作。这仍然是O(log )，但它可能会稍微快一点。

Answer 2

如果您有一个迭代器，并且不需要在异步中执行，则可以使用std::partial_sort。

std::vector<int> x{10, 5, 1, 2, 3};
std::partial_sort(x.begin(), x.begin() + k, x.end());

这使得前k个元素的运行时复杂度大致为O(nlogk)。