BLI_inplace_priority_queue.hh

Go to the documentation of this file.

/* SPDX-FileCopyrightText: 2023 Blender Authors
 *
 * SPDX-License-Identifier: GPL-2.0-or-later */
 
#pragma once
 
#include "BLI_array.hh"
#include "BLI_dot_export.hh"
 
#include <sstream>
 
namespace blender {
 
template<
    /* Type of the elements in the underlying array. */
    typename T,
    /* Binary function that takes two `const T &` inputs and returns true,
     * when the first input has greater priority than the second. */
    typename FirstHasHigherPriority = std::greater<T>>
class InplacePriorityQueue {
 private:
  /* Underlying array the priority queue is built upon. This is a span instead of a mutable span,
   * because this data structure never changes the values itself. */
  Span<T> data_;
  /* Maps indices from the heap (binary tree in array format) to indices of the underlying/original
   * array. */
  Array<int64_t> heap_to_orig_;
  /* This is the inversion of the above mapping. */
  Array<int64_t> orig_to_heap_;
  /* Number of elements that are currently in the priority queue. */
  int64_t heap_size_ = 0;
  /* Function that can be changed to customize how the priority of two elements is compared. */
  FirstHasHigherPriority first_has_higher_priority_fn_;
 
 public:
  InplacePriorityQueue(Span<T> data)
      : data_(data), heap_to_orig_(data_.size()), orig_to_heap_(data_.size())
  {
    for (const int64_t i : IndexRange(data_.size())) {
      heap_to_orig_[i] = i;
      orig_to_heap_[i] = i;
    }
 
    this->rebuild();
  }
 
  void rebuild()
  {
    const int final_heap_size = data_.size();
    if (final_heap_size > 1) {
      for (int64_t i = this->get_parent(final_heap_size - 1); i >= 0; i--) {
        this->heapify(i, final_heap_size);
      }
    }
    heap_size_ = final_heap_size;
  }
 
  int64_t size() const
  {
    return heap_size_;
  }
 
  bool is_empty() const
  {
    return heap_size_ == 0;
  }
 
  const T &peek() const
  {
    return data_[this->peek_index()];
  }
 
  const T &pop()
  {
    return data_[this->pop_index()];
  }
 
  int64_t peek_index() const
  {
    BLI_assert(!this->is_empty());
    return heap_to_orig_[0];
  }
 
  int64_t pop_index()
  {
    BLI_assert(!this->is_empty());
    const int64_t top_index_orig = heap_to_orig_[0];
    heap_size_--;
    if (heap_size_ > 1) {
      this->swap_indices(0, heap_size_);
      this->heapify(0, heap_size_);
    }
    return top_index_orig;
  }
 
  void priority_decreased(const int64_t index)
  {
    const int64_t heap_index = orig_to_heap_[index];
    if (heap_index >= heap_size_) {
      /* This element is not in the queue currently. */
      return;
    }
    this->heapify(heap_index, heap_size_);
  }
 
  void priority_increased(const int64_t index)
  {
    int64_t current = orig_to_heap_[index];
    if (current >= heap_size_) {
      /* This element is not in the queue currently. */
      return;
    }
    while (true) {
      if (current == 0) {
        break;
      }
      const int64_t parent = this->get_parent(current);
      if (this->first_has_higher_priority(parent, current)) {
        break;
      }
      this->swap_indices(current, parent);
      current = parent;
    }
  }
 
  void priority_changed(const int64_t index)
  {
    this->priority_increased(index);
    this->priority_decreased(index);
  }
 
  Span<int64_t> active_indices() const
  {
    return heap_to_orig_.as_span().take_front(heap_size_);
  }
 
  Span<int64_t> inactive_indices() const
  {
    return heap_to_orig_.as_span().drop_front(heap_size_);
  }
 
  Span<int64_t> all_indices() const
  {
    return heap_to_orig_;
  }
 
  std::string to_dot() const
  {
    return this->partial_to_dot(heap_size_);
  }
 
 private:
  bool first_has_higher_priority(const int64_t a, const int64_t b)
  {
    const T &value_a = data_[heap_to_orig_[a]];
    const T &value_b = data_[heap_to_orig_[b]];
    return first_has_higher_priority_fn_(value_a, value_b);
  }
 
  void swap_indices(const int64_t a, const int64_t b)
  {
    std::swap(heap_to_orig_[a], heap_to_orig_[b]);
    orig_to_heap_[heap_to_orig_[a]] = a;
    orig_to_heap_[heap_to_orig_[b]] = b;
  }
 
  void heapify(const int64_t parent, const int64_t heap_size)
  {
    int64_t max_index = parent;
    const int left = this->get_left(parent);
    const int right = this->get_right(parent);
    if (left < heap_size && this->first_has_higher_priority(left, max_index)) {
      max_index = left;
    }
    if (right < heap_size && this->first_has_higher_priority(right, max_index)) {
      max_index = right;
    }
    if (max_index != parent) {
      this->swap_indices(parent, max_index);
      this->heapify(max_index, heap_size);
    }
    if (left < heap_size) {
      BLI_assert(!this->first_has_higher_priority(left, parent));
    }
    if (right < heap_size) {
      BLI_assert(!this->first_has_higher_priority(right, parent));
    }
  }
 
  int64_t get_parent(const int64_t child) const
  {
    BLI_assert(child > 0);
    return (child - 1) / 2;
  }
 
  int64_t get_left(const int64_t parent) const
  {
    return parent * 2 + 1;
  }
 
  int64_t get_right(const int64_t parent) const
  {
    return parent * 2 + 2;
  }
 
  std::string partial_to_dot(const int size) const
  {
    dot::DirectedGraph digraph;
    Array<dot::Node *> dot_nodes(size);
    for (const int i : IndexRange(size)) {
      std::stringstream ss;
      ss << data_[heap_to_orig_[i]];
      const std::string name = ss.str();
      dot::Node &node = digraph.new_node(name);
      node.set_shape(dot::Attr_shape::Rectangle);
      node.attributes.set("ordering", "out");
      dot_nodes[i] = &node;
      if (i > 0) {
        const int64_t parent = this->get_parent(i);
        digraph.new_edge(*dot_nodes[parent], node);
      }
    }
    return digraph.to_dot_string();
  }
};
 
}  // namespace blender