std::priority_queue
does not support dynamically updating element priorities. This class template extends it by allowing updating priorities of elements in the heap. Any comments are welcome. In particular, I hope that the implementation is at least bug-free. Original code and unit test can be found here and here.
#include <cassert>
#include <cstddef>
#include <algorithm>
#include <functional>
#include <initializer_list>
#include <unordered_map>
#include <utility>
#include <vector>
/// Binary max-heap with mutable priorities. `Less` is the less-than comparator
/// type that compares the priorities of objects referenced by key type
/// `std::intptr_t`.
template <class Less = std::less<std::intptr_t>>
class mutable_heap {
public:
using key_type = std::intptr_t;
using pos_type = std::size_t;
using map_type = std::unordered_map<key_type, pos_type>;
using vec_type = std::vector<key_type>;
explicit mutable_heap(Less less = Less{})
: less_(std::move(less)) {
// nop
}
template <class ForwardIt>
mutable_heap(ForwardIt first, ForwardIt last, Less less = Less{})
: less_(std::move(less)),
heap_(first, last) {
std::make_heap(heap_.begin(), heap_.end(), less_);
for (pos_type i = 0; i < heap_.size(); ++i) {
key_pos_map_.emplace(heap_[i], i);
}
}
explicit mutable_heap(std::initializer_list<std::intptr_t> list,
Less less = Less{})
: mutable_heap(list.begin(), list.end(), std::move(less)) {
// nop
}
key_type top() const {
assert(! empty());
return heap_.front();
}
key_type pop() {
assert(! empty());
auto key = top();
heap_.front() = heap_.back();
auto& key_pos = key_pos_map_[heap_.front()] = 0;
key_pos_map_.erase(key);
heap_.pop_back();
sift_down(heap_.begin(), heap_.end(), key_pos);
return key;
}
/// Add the object referenced by `key` to the heap. If it's already in the
/// heap, its position in the heap is updated if necessary.
void push(key_type key) {
auto pr = key_pos_map_.emplace(key, heap_.size());
if (pr.second) {
heap_.emplace_back(key);
sift_up(heap_.begin(), heap_.end(), pr.first->second);
return;
}
auto it = heap_.begin() + pr.first->second;
if (sift_up(heap_.begin(), it + 1, pr.first->second)) return;
sift_down(it, heap_.end(), pr.first->second);
return;
}
void emplace(key_type key) {
push(key);
}
void clear() {
heap_.clear();
key_pos_map_.clear();
}
auto size() const {
return heap_.size();
}
bool empty() const {
return heap_.empty();
}
void reserve(pos_type cap) {
heap_.reserve(cap);
key_pos_map_.reserve(cap);
}
private:
bool sift_up(vec_type::iterator first, vec_type::iterator last,
pos_type& key_pos) {
auto n = static_cast<pos_type>(last - first);
if (n < 2) return false;
auto parent = n / 2 - 1;
if (! less_(first[parent], last[-1])) return false;
auto key = std::exchange(last[-1], first[parent]);
key_pos_map_[last[-1]] = static_cast<pos_type>(last - first - 1);
while (true) {
last = first + parent + 1;
n = static_cast<pos_type>(last - first);
if (n < 2) break;
parent = n / 2 - 1;
if (! less_(first[parent], key)) break;
last[-1] = first[parent];
key_pos_map_[last[-1]] = static_cast<pos_type>(last - first - 1);
}
last[-1] = key;
key_pos = static_cast<pos_type>(last - first - 1);
return true;
}
bool sift_down(vec_type::iterator first, vec_type::iterator last,
pos_type& key_pos) {
auto first_pos = static_cast<pos_type>(first - heap_.begin());
auto last_pos = static_cast<pos_type>(last - heap_.begin());
auto child_pos = 2 * first_pos + 1;
if (! (child_pos < last_pos)) return false;
if (child_pos + 1 < last_pos
&& less_(heap_[child_pos], heap_[child_pos + 1])) {
++child_pos;
}
if (! less_(heap_[first_pos], heap_[child_pos])) return false;
auto key = std::exchange(heap_[first_pos], heap_[child_pos]);
key_pos_map_[heap_[first_pos]] = first_pos;
while (true) {
first_pos = child_pos;
child_pos = first_pos * 2 + 1;
if (! (child_pos < last_pos)) break;
if (child_pos + 1 < last_pos
&& less_(heap_[child_pos], heap_[child_pos + 1])) {
++child_pos;
}
if (! less_(heap_[first_pos], heap_[child_pos])) break;
heap_[first_pos] = heap_[child_pos];
key_pos_map_[heap_[first_pos]] = first_pos;
}
heap_[first_pos] = key;
key_pos = first_pos;
return true;
}
Less less_;
vec_type heap_;
map_type key_pos_map_;
};
std::intptr_t
. So it is generic, kind of. \$\endgroup\$ – Lingxi Feb 6 '16 at 2:00