heap.hpp

#ifndef CUSTOM_HEAP_HPP
#define CUSTOM_HEAP_HPP

#include <vector>
#include <cstddef>

// Binary max-heap.. generic.. uses a KeyExtractor functor
// push, top, pop, empty, size .. implemented

template<typename Item, typename KeyExtractor>
class BinaryMaxHeap {
private:
    std::vector<Item> a;
    KeyExtractor getKey;

    inline size_t parent(size_t i) const { return (i - 1) / 2; }
    inline size_t left(size_t i) const { return 2 * i + 1; }
    inline size_t right(size_t i) const { return 2 * i + 2; }

    bool cmp(size_t i, size_t j) const {
        return getKey(a[i]) > getKey(a[j]);
    }

    void shiftUp(size_t i) {
        while (i > 0 && !cmp(parent(i), i)) {
            std::swap(a[i], a[parent(i)]);
            i = parent(i);
        }
    }

    void shiftDown(size_t i) {
        while (true) {
            size_t l = left(i), r = right(i);
            size_t best = i;
            if (l < a.size() && !cmp(best, l)) best = l;
            if (r < a.size() && !cmp(best, r)) best = r;
            if (best == i) break;
            std::swap(a[i], a[best]);
            i = best;
        }
    }

public:
    BinaryMaxHeap() : a(), getKey(KeyExtractor()) {}

    void push(const Item& it) {
        a.push_back(it);
        shiftUp(a.size() - 1);
    }

    const Item& top() const {
        return a.front();
    }

    void pop() {
        std::swap(a.front(), a.back());
        a.pop_back();
        if (!a.empty()) shiftDown(0);
    }

    bool empty() const { return a.empty(); }
    size_t size() const { return a.size(); }
};

#endif