From acd6a581e7764a38a5062c50f671ed259957a8cc Mon Sep 17 00:00:00 2001 From: Forrest Green Date: Wed, 22 Mar 2017 09:17:39 -0700 Subject: [PATCH 1/2] Add RankedTree. I didn't end up using this, but it's fairly complete and might be useful later --- geode/structure/CMakeLists.txt | 3 + geode/structure/RankedTree.cpp | 103 ++++++ geode/structure/RankedTree.h | 229 ++++++++++++++ geode/structure/RankedTree_p.h | 475 ++++++++++++++++++++++++++++ geode/structure/module.cpp | 1 + geode/structure/test_ranked_tree.py | 8 + 6 files changed, 819 insertions(+) create mode 100644 geode/structure/RankedTree.cpp create mode 100644 geode/structure/RankedTree.h create mode 100644 geode/structure/RankedTree_p.h create mode 100644 geode/structure/test_ranked_tree.py diff --git a/geode/structure/CMakeLists.txt b/geode/structure/CMakeLists.txt index 046fd3f5..b597fca9 100644 --- a/geode/structure/CMakeLists.txt +++ b/geode/structure/CMakeLists.txt @@ -1,5 +1,6 @@ set(module_SRCS Heap.cpp + RankedTree.cpp Tuple.cpp ) @@ -13,6 +14,8 @@ set(module_HEADERS Quad.h Queue.h Quintuple.h + RankedTree.h + RankedTree_p.h Singleton.h Stack.h Triple.h diff --git a/geode/structure/RankedTree.cpp b/geode/structure/RankedTree.cpp new file mode 100644 index 00000000..b0a733ce --- /dev/null +++ b/geode/structure/RankedTree.cpp @@ -0,0 +1,103 @@ +#include "RankedTree.h" +#include +#include +#include +#include + +namespace geode { + + +static auto random_permutation(const int length, const int seed) -> Array { + uint128_t key = 12345 + seed; + const auto result = Array{length, uninit}; + for(const int i : range(length)) { + result[i] = static_cast(random_permute(length, key, i)); + } + return result; +} + +// Add a bunch of integers to a tree in order making sure that that is what tree contains at every step +static void try_integer_range() { + for(const int seed : {0,1,2,3}) { + for(const int test_n : {1,2,3,4,5,10,15,16,30,31,32,33,34,100,200,300}) { + const Array test_data = random_permutation(test_n, seed + test_n*4); + RankedTree l; + GEODE_ASSERT(l.empty()); // New tree should be empty + l.test_global_invarients(); + for(const int i : range(test_n)) { + l.emplace_back(test_data[i]); + l.test_global_invarients(); + int expected = 0; + for(const int j : l) { + GEODE_ASSERT(j == test_data[expected++]); + } + GEODE_ASSERT(expected == i+1); + } + + // Check that we can find values by index + for(const int i : range(test_n)) { + const auto iter = l.find_index(i); + GEODE_ASSERT(iter != l.end() && *iter == test_data[i]); + } + + // Now go back and remove things one at a time + for(const auto rev_i : range(1,test_n+1)) { + const int i = test_data[test_n - rev_i]; + GEODE_ASSERT(!l.empty()); + const auto iter = l.find_last(); + GEODE_ASSERT(*iter == i); + GEODE_ASSERT(l.erase(iter) == l.end()); + l.test_global_invarients(); + int expected = 0; + for(const int j : l) { + GEODE_ASSERT(j == test_data[expected++]); + } + GEODE_ASSERT(expected == test_data.size() - rev_i); + } + GEODE_ASSERT(l.empty()); + } + } +} + +static void try_insertion_sort() { + constexpr int test_length = 50; + const auto test_data = random_permutation(test_length, 0); + RankedTree l; + const auto is_sorted = [](RankedTree& l) { + int prev = -1; + for(const int i : l) { + if(!(prev < i)) return false; + prev = i; + } + return true; + }; + + for(const int i : range(test_length)) { + const int n = test_data[i]; + int n_evaluations = 0; + const auto p = [&](int e) { n_evaluations += 1; return e < n; }; + l.emplace_in_order(p, n); + GEODE_ASSERT(n_evaluations <= (integer_log(i) + 2)); + GEODE_ASSERT(is_sorted(l)); + l.test_global_invarients(); + } + int expected = 0; + for(const int i : l) { + GEODE_ASSERT(i == expected); + expected += 1; + } + GEODE_ASSERT(expected == test_length); +} + +void ranked_tree_test() { + try_integer_range(); + try_insertion_sort(); +} + +} // geode namespace + +using namespace geode; + +void wrap_ranked_tree() { + GEODE_FUNCTION(ranked_tree_test) +} diff --git a/geode/structure/RankedTree.h b/geode/structure/RankedTree.h new file mode 100644 index 00000000..af6f26f9 --- /dev/null +++ b/geode/structure/RankedTree.h @@ -0,0 +1,229 @@ +#pragma once +#include +#include +#include +#include +#include +#include + +namespace geode { + +// Stores a collection of elements in a fixed order with support for the following operation: +// Insertion, and removal of elements from anywhere in log time +// Reorder any two elements in constant time +// Binary search over the entire collection (assuming it is sorted) in log time +// Iterate over consecutive elements in amortized constant time +// Find an element at a specific index in log time +// This class is intended to store a sorted collection updating it as elements are added and removed or as sort order changes +// Unlike std::set/map/multiset/multimap this class does not maintain ordered keys for entries, allowing for a comparison function that might change +// (Though the assumption is that any changes to comparison function will only alter a small number of nodes) +// Caller is responsible for updating nodes when their order changes for binary search to work +template class RankedTree { + private: + class Node; + + class IteratorBase { + protected: + friend class RankedTree; + // Traversal order should be left, self, right + Node* node_ptr = nullptr; + Node& node() const { assert(node_ptr); return *node_ptr; } + explicit IteratorBase(Node* n) : node_ptr(n) { } + public: + IteratorBase() = default; + explicit operator bool() const { return node_ptr != nullptr; } + inline friend bool operator==(const IteratorBase& lhs, const IteratorBase& rhs) { return lhs.node_ptr == rhs.node_ptr; } + inline friend bool operator!=(const IteratorBase& lhs, const IteratorBase& rhs) { return lhs.node_ptr != rhs.node_ptr; } + }; + public: + class ConstIterator : public IteratorBase { + protected: + friend class RankedTree; + friend class Iterator; + explicit ConstIterator(Node* n) : IteratorBase(n) { } + using IteratorBase::node_ptr; + public: + ConstIterator& operator--() { assert(node_ptr); node_ptr = node_ptr->prev(); return *this; } + ConstIterator& operator++() { assert(node_ptr); node_ptr = node_ptr->next(); return *this; } + ConstIterator prev() const { assert(node_ptr); return ConstIterator{node_ptr->prev()}; } + ConstIterator next() const { assert(node_ptr); return ConstIterator{node_ptr->next()}; } + const T& operator*() const { assert(node_ptr); return node_ptr->data; } + }; + class Iterator : public IteratorBase { + protected: + friend class RankedTree; + explicit Iterator(Node* n) : IteratorBase(n) { } + using IteratorBase::node_ptr; + public: + // Allow implicit conversion to ConstIterator + operator ConstIterator() const { return ConstIterator{node_ptr}; } + Iterator& operator--() { assert(node_ptr); node_ptr = node_ptr->prev(); return *this; } + Iterator& operator++() { assert(node_ptr); node_ptr = node_ptr->next(); return *this; } + Iterator prev() const { assert(node_ptr); return Iterator{node_ptr->prev()}; } + Iterator next() const { assert(node_ptr); return Iterator{node_ptr->next()}; } + T& operator*() const { assert(node_ptr); return node_ptr->data; } + }; + + + bool empty() const; + Iterator begin(); + ConstIterator begin() const; + Iterator end(); + ConstIterator end() const; + + Iterator find_last(); + ConstIterator find_last() const; + + // Inserts a new element directly before pos, returns iterator to newly inserted element. O(log(size)) + // No iterators or references are invalidated + template Iterator emplace(const ConstIterator pos, Args&&... args) + { return insert_before(pos, NodeLink{new Node{std::forward(args)...}}); } + template Iterator emplace_after(const ConstIterator pos, Args&&... args) + { return insert_after(pos, NodeLink{new Node{std::forward(args)...}}); } + template Iterator emplace_back(Args&&... args) + { return insert_before(end(), NodeLink{new Node{std::forward(args)...}}); } + + // Removes element at pos. Returns iterator to next element after removed iterator + // Only iterators pointing to erased element should be invalidated + Iterator erase(const ConstIterator pos); + + // Swaps position within list of two elements + // No iterators are invalidated. Addresses do not change + void swap_positions(const ConstIterator pos0, const ConstIterator pos1) { assert(pos0 && pos1); swap_positions(pos0.node(), pos1.node()); } + + // Locate element at index i. O(log(size)) + Iterator find_index(size_t i); + ConstIterator find_index(size_t i) const; + + // The following functions assume element are partially ordered with respect to searched element or predicate + // i.e. Order should be as if std::partition has been called with predicate + + // Finds first element for which p returns false (p) + // Returns end if no such element exists + template ConstIterator find_first_false(const UnaryPredicate& p) const; + template Iterator find_first_false(const UnaryPredicate& p) + { return Iterator{static_cast(this)->find_first_false(p).node_ptr}; } + // Assuming p returns <0 for elements before range, ==0 for elements inside range, and >0 for elements after range + // Finds [start, end) containing all elements that return ==0 + template Range equal_range(const SignedPredicate& p) const; + + + // Inserts a new element before first element for which is_before(element) is false + template Iterator emplace_in_order(const UnaryPredicate& is_before, Args&&... args) + { return insert_before(find_first_false(is_before), NodeLink{new Node{std::forward(args)...}}); } + + // Dump internal data for debugging + void debug_print(std::ostream& os) const { debug_print(os, '_', 0, m_root.get()); } + + // Verifies expected properties for every node + void test_global_invarients() const; + private: + // Implementation is based on a left leaning red black tree with a 'rank' added to each node that can be converted to an index during iteration + // 'rank' tracks total number of nodes in left subtree which is updated as part of any re-balancing operations + // Ownership of nodes is managed via std::unique_ptr in parent node or at root of tree + + // Enums for tracking properties of + enum class Color : bool { RED, BLACK }; + enum class Side : bool { LEFT, RIGHT }; + inline friend Color operator!(Color c) { return static_cast(!static_cast(c)); } + inline friend Side operator^(Side s, bool b) { return static_cast(static_cast(s) ^ b); } + using NodeLink = std::unique_ptr; + static bool is_red(Node* n) { return n && n->is_red(); } + static bool is_black(Node* n) { return !n || n->is_black(); } + static bool is_red(const NodeLink& n) { return is_red(n.get()); } + static bool is_black(const NodeLink& n) { return is_black(n.get()); } + class Node { + public: + NodeLink left; // Left nodes come before this + NodeLink right; // Right nodes come after this + Node* parent = nullptr; + static constexpr size_t combine(size_t rank, Color c) { return (rank<<1)|static_cast(c); } + size_t m_rank_and_color = combine(1, Color::RED); + T data; + const T& cdata() const { return data; } // Shorthand for casting data to const + // Constructor forwards all arguments to T + template explicit Node(Args&&... args) + : data(std::forward(args)...) + { } + ~Node() { assert(!parent || (parent->left.get() != this && parent->right.get() != this)); } + + size_t rank() const { return m_rank_and_color>>1; } + void update_rank(ssize_t delta) { assert(rank() + delta > 0); m_rank_and_color += (delta<<1); } + Color color() const { return static_cast(m_rank_and_color & 1); } + void set_color(const Color new_color) { m_rank_and_color = combine(rank(), new_color); } + void set_rank(const size_t new_rank) { m_rank_and_color = combine(new_rank, color()); } + + bool is_red() const { return color() == Color::RED; } + bool is_black() const { return color() == Color::BLACK; } + static bool is_red(Node* n) { return RankedTree::is_red(n); } + static bool is_black(Node* n) { return RankedTree::is_black(n); } + + bool is_root() const { return !parent; } + bool is_left() const { assert(parent && ((parent->left.get() == this) != (parent->right.get() == this))); return (parent->left.get() == this); } + bool is_right() const { assert(parent && ((parent->left.get() == this) != (parent->right.get() == this))); return (parent->right.get() == this); } + + Side side() const; + NodeLink link(const Side s) { return (s == Side::LEFT) ? left : right; } + + Node& min_child() { return left ? left->min_child() : *this; } + Node& max_child() { return right ? right->max_child() : *this; } + + Node* next(); // Find first node after this or null if this node is the max + Node* prev(); // Previous node before this or null if this node is the min + + // Validates expected invariants for an individual node and connections to neighbors + void test_local_invarients() const; + }; + + NodeLink m_root; + + // Returns the link that owns node. This will be either m_root, node.parent->left, or node.parent->right + NodeLink& parent_link(Node& node); + + // Swaps all auxiliary data and references to/from n0/n1 but leaves Node::data unchanged + // This could almost be replaced with swap(n0.data, n1.data), but that invalidates iterators and requires T is swappable + void swap_positions(Node& n0, Node& n1); + + // Swaps h and h->right adjusting other members accordingly + static void rotate_left(NodeLink& h); + // Swaps h and h->left adjusting other members accordingly + static void rotate_right(NodeLink& h); + // Flips color of h and children of h + static void flip_color(Node& node); + // Rotates h as needed to maintain invariants after inserting or deleting a node + static bool fixup(NodeLink& h); + // Makes node red without unbalancing tree. Might leave ancestors right-leaning + void force_red(Node& node); + // Walks from h to root of tree handling needed changes after inserting or deleting a node + void fixup_ancestors(NodeLink& h); + + // Moves new_node into h then walks up tree re-balancing nodes as needed + // Requires h is a currently empty leaf of the tree + void link_node(Node* parent, NodeLink& h, NodeLink&& new_node); + + // Helper functions that find suitable parent node and then call link_node + Iterator insert_before(const Iterator pos, NodeLink&& new_element); + Iterator insert_after(const Iterator pos, NodeLink&& new_element); + + // Removes node from the tree, reattaching children in the tree in the correct order and re-balancing as needed + // Node must be valid when calling the function but will be deleted by the time this function returns + void extract_node(Node* node); + + struct SubtreeStats { + size_t black_height; // Total number of black links in this subtree (should be the same for any path to a leaf) + size_t total_size; // Total count of nodes in this subtree + size_t max_height; // Longest path to a leaf from root of this subtree + }; + // + SubtreeStats test_subtree_invarients(const Node* node) const; + + void debug_print(std::ostream& os, const char prefix, const int indent, const Node * node) const; +}; + +// This printing operator includes a bunch of debug information and might not be ideal for general use +template std::ostream& operator<<(std::ostream& os, const RankedTree& tree) +{ tree.debug_print(os); return os; } + +} // geode namespace + +#include "RankedTree_p.h" diff --git a/geode/structure/RankedTree_p.h b/geode/structure/RankedTree_p.h new file mode 100644 index 00000000..7798ae97 --- /dev/null +++ b/geode/structure/RankedTree_p.h @@ -0,0 +1,475 @@ +// Implementation of template methods for RankedTree +// This should only be included from RankedTree.h +#pragma once +namespace geode { + +template auto RankedTree::debug_print(std::ostream& os, const char prefix, const int indent, const Node* node) const -> void { + for(int i = 0; i < indent; ++i) os << ' '; + os << prefix << ' '; + if(node) { + os << (node->is_red() ? "RED " : "BLACK ") << node->data << " @ " << node << " of rank " << node->rank() << '\n'; + debug_print(os, 'L', indent + 2, node->left.get()); + debug_print(os, 'R', indent + 2, node->right.get()); + } + else { + os << "BLACK leaf @ " << node << '\n'; + } +} + +template auto RankedTree::Node::test_local_invarients() const -> void { + GEODE_ASSERT(parent != this); // Parent should never point back to self (Nor should cycles exist anywhere, but this an easy one to check) + GEODE_ASSERT(parent || is_black()); // Root node is always black + GEODE_ASSERT(!parent || parent->left.get() == this || parent->right.get() == this); // If parent is set, this should be one of the children + GEODE_ASSERT(!left || left->parent == this); // Children should point back to self + GEODE_ASSERT(!right || right->parent == this); // Children should point back to self + GEODE_ASSERT(is_black() || (is_black(parent) && is_black(left.get()) && is_black(right.get()))); // No consecutive red nodes + GEODE_ASSERT(left ? true : !right); // Node should always be left leaning + GEODE_ASSERT(!(is_red(left.get()) && is_red(right.get()))); // Should never have 4 nodes + GEODE_ASSERT(rank() >= 1 + (left ? left->rank() : 0)); +} + + +template auto RankedTree::Node::side() const -> Side { + assert(parent); // Should probably not be calling this function on the root + if(parent) { + assert((this == parent->left) != (this == parent->right)); + return (this == parent->left) ? Side::LEFT : Side::RIGHT; + } + else { + return Side::LEFT; + } +} + +template auto RankedTree::Node::next() -> Node* { + // Next node is min below right if it exists + if(right) return &right->min_child(); + // Walk upwards until we find a node with a parent greater than where we started + for(Node* iter = this; iter->parent; iter = iter->parent) { + if(iter->is_left()) + return iter->parent; + } + return nullptr; // If every ancestor was less current node is max +} + +template auto RankedTree::Node::prev() -> Node* { + // Logic is as Node::next with directions reversed + if(left) return &left->max_child(); + for(Node* iter = this; iter->parent; iter = iter->parent) { + if(iter->is_right()) + return iter->parent; + } + return nullptr; +} + +template auto RankedTree::parent_link(Node& node) -> NodeLink& { + if(node.parent) { + return node.is_left() ? node.parent->left : node.parent->right; + } + else { + assert(m_root.get() == &node); + return m_root; // Link to root node is m_root + } +} + +template auto RankedTree::test_subtree_invarients(const Node* node) const -> SubtreeStats { + SubtreeStats result; + if(!node) { + result.total_size = 0; + result.black_height = 1; + result.max_height = 1; + } + else { + // Only the root node should have no parent + GEODE_ASSERT((node->parent == nullptr) == (node == m_root.get())); + node->test_local_invarients(); + + const auto l_stats = test_subtree_invarients(node->left.get()); + const auto r_stats = test_subtree_invarients(node->right.get()); + + GEODE_ASSERT(node->rank() == l_stats.total_size + 1); + result.total_size = l_stats.total_size + r_stats.total_size + 1; + GEODE_ASSERT(l_stats.black_height == r_stats.black_height); + result.black_height = l_stats.black_height; // Use left value since they are the same + if(node->is_black()) { + result.black_height += 1; + } + result.max_height = 1 + max(l_stats.max_height, r_stats.max_height); + GEODE_ASSERT(2*result.black_height >= result.max_height); // Should have at least every other link black + GEODE_ASSERT(result.total_size <= 1< auto RankedTree::test_global_invarients() const -> void +{ test_subtree_invarients(m_root.get()); } + +// This shouldn't alter any of the properties of the graph +template auto RankedTree::swap_positions(Node& n0, Node& n1) -> void { + assert(&n0 != &n1); + if(&n0 == &n1) return; + std::swap(parent_link(n0), parent_link(n1)); + std::swap(n0.left, n1.left); + std::swap(n0.right, n1.right); + std::swap(n0.parent, n1.parent); + std::swap(n0.m_rank_and_color, n1.m_rank_and_color); + // Fix parent points on children + for(Node* n : {&n0, &n1}) { + assert(n->left || !n->right); // Check node is left leaning so we don't need to check right link + if(n->left) { + assert(n->left->parent != n); // Should be pointing to other node we are swapping with + n->left->parent = n; + if(n->right) { + assert(n->right->parent != n); // Should be pointing to other node we are swapping with + n->right->parent = n; + } + } + } +} + +// Swaps h and h->right adjusting colors and other pointers accordingly +template auto RankedTree::rotate_left(NodeLink& h) -> void { + assert(h && h->right); + h->right->set_color(h->color()); + h->set_color(Color::RED); + auto& orig_h = *h; + auto& orig_x = *(h->right); + swap(h, orig_h.right); + swap(orig_x.left, orig_h.right); + orig_x.parent = orig_h.parent; + orig_h.parent = &orig_x; + if(orig_h.right) { + assert(orig_h.right->parent == &orig_x); + orig_h.right->parent = &orig_h; + } + orig_x.update_rank(orig_h.rank()); + assert(h.get() == &orig_x); + assert(h->left.get() == &orig_h); +} + +// Swaps h and h->right adjusting colors and other pointers accordingly +template auto RankedTree::rotate_right(NodeLink& h) -> void { + assert(h && h->left); + h->left->set_color(h->color()); + h->set_color(Color::RED); + auto& orig_h = *h; + auto& orig_x = *(h->left); + swap(h, orig_h.left); + swap(orig_x.right, orig_h.left); + orig_x.parent = orig_h.parent; + orig_h.parent = &orig_x; + if(orig_h.left) { + assert(orig_h.left->parent == &orig_x); + orig_h.left->parent = &orig_h; + } + orig_h.update_rank(-orig_x.rank()); + assert(h.get() == &orig_x); + assert(h->right.get() == &orig_h); +} + +template auto RankedTree::flip_color(Node& node) -> void { + assert(node.left && node.right); + assert(node.left->color() == node.right->color()); + node.m_rank_and_color ^= 1; + node.left->m_rank_and_color ^= 1; + node.right->m_rank_and_color ^= 1; +} + +template auto RankedTree::fixup(NodeLink& h) -> bool { + bool changed = false; + // Fix any newly created right-leaning nodes + if(is_red(h->right) && !is_red(h->left)) { + changed = true; + rotate_left(h); + } + + // Fix two red nodes in a row + if(is_red(h->left) && is_red(h->left->left)) { + changed = true; + rotate_right(h); + } + + // Split any 4 nodes + if(is_red(h->left) && is_red(h->right)) { + changed = true; + flip_color(*h); + } + + if(!h->parent) h->set_color(Color::BLACK); // Reset root node to black + + return changed; +} + +template auto RankedTree::fixup_ancestors(NodeLink& mutated) -> void { + NodeLink* iter = &mutated; + int consecutive_unchanged_links = fixup(*iter) ? 0 : 1; + while(iter != &m_root) { + iter = &parent_link(*(*iter)->parent); + if(fixup(*iter)) { + consecutive_unchanged_links = 0; + } + else { + // Properties that fixup is trying to repair depend on 3 consecutive levels of tree + // If we didn't change this level or the previous two, we don't need to worry about other ancestors + if(consecutive_unchanged_links >= 2) + return; + ++consecutive_unchanged_links; + } + } +} + +// Sets h to new_node and then re-balances the tree +template auto RankedTree::link_node(Node* parent, NodeLink& h, NodeLink&& new_node) -> void { + assert(!h); // h should be empty + assert(parent ? (&h == &parent->left || &h == &parent->right) + : (&h == &m_root)); // h should be a link from the indicated parent (or the root link) + assert(new_node); // We expect to be getting a non-null node to insert + assert(!new_node->left && !new_node->right); // This function assumes new_node doesn't already have children which might unbalance tree more than could be fixed here + assert(!new_node->parent); // We assume new_node doesn't already have some parent it needs to be detached from + assert(new_node->rank() == 1); // Rank should be initialized to one + assert(new_node->is_red()); // New nodes are red + h = std::move(new_node); + h->parent = parent; + assert(h->parent != h.get()); // Check that we didn't try to make a node it's own parent + // Update ranks all the way back to the root + // TODO: It would probably be faster to combine this with fixup_ancestors + for(Node* iter = h.get(); iter->parent; iter = iter->parent) { + if(iter->is_left()) iter->parent->update_rank(1); + } + + // We've attached a new red node so bubble up any fixes we need to make to avoid double reds or other constraints + fixup_ancestors(h); +} + +template auto RankedTree::insert_after(const Iterator pos, NodeLink&& new_node) -> Iterator { + const auto result = Iterator{new_node.get()}; + assert(pos); // Shouldn't try to insert after end + assert(!empty()); // If pos is valid and this is empty, something went wrong (maybe pos is from a different tree?) + + // Inserting after end isn't intentionally supported, but seems safer to stuff value at end instead of dereferencing a null pointer + if(!pos) return insert_before(end(), std::move(new_node)); + + // Find node with empty link to attach to + Node* iter = pos.node_ptr; + while(iter->right) { + assert(iter->left); // Left leaning + iter = iter->left.get(); + } + + // Attaching on right temporarily violates left leaning property, but link_node will fix that + link_node(iter, iter->right, std::move(new_node)); + assert(pos.node_ptr->next() == result.node_ptr); + assert(result.node_ptr->prev() == pos.node_ptr); + return result; +} + +template auto RankedTree::insert_before(const Iterator pos, NodeLink&& new_node) -> Iterator { + if(empty()) { // Special case handling for empty tree + assert(!pos); // It shouldn't be possible to have a valid iterator into an empty tree + m_root = std::move(new_node); + m_root->set_color(Color::BLACK); + return Iterator{m_root.get()}; + } + const auto result = Iterator{new_node.get()}; + + Node* iter = pos.node_ptr; + + if(!iter) { // Special case handling for end of tree + iter = &(m_root->max_child()); + assert(iter); // Should have caught empty tree above + // Use insert_after to find correct spot + return insert_after(Iterator{iter}, std::move(new_node)); + } + + if(!iter->left) { + link_node(iter, iter->left, std::move(new_node)); + } + else { + iter = &(iter->left->max_child()); + // Attaching on right temporarily violates left leaning property, but link_node will fix that + link_node(iter, iter->right, std::move(new_node)); + } + assert(result.node_ptr->next() == pos.node_ptr); + return result; +} + +template auto RankedTree::find_last() -> Iterator { + if(!m_root) return Iterator{nullptr}; + else return Iterator{&(m_root->max_child())}; +} + +template auto RankedTree::find_last() const -> ConstIterator { + if(!m_root) return ConstIterator{nullptr}; + else return ConstIterator{&(m_root->max_child())}; +} + +template auto RankedTree::empty() const -> bool { return m_root == nullptr; } +template auto RankedTree::begin() const -> ConstIterator { return ConstIterator{m_root ? &m_root->min_child() : nullptr}; } +template auto RankedTree::begin() -> Iterator { return Iterator{m_root ? &m_root->min_child() : nullptr}; } +template auto RankedTree::end() const -> ConstIterator { return Iterator{nullptr}; } +template auto RankedTree::end() -> Iterator { return Iterator{nullptr}; } + +template auto RankedTree::erase(ConstIterator pos) -> Iterator { + assert(pos); + const auto result = Iterator{pos.node().next()}; + extract_node(pos.node_ptr); + return result; +} + +template auto RankedTree::find_index(const size_t i) -> Iterator { + size_t m = i + 1; + Node* iter = m_root.get(); + while(iter) { + if(m < iter->rank()) { + iter = iter->left.get(); + } + else if(m > iter->rank()) { + m -= iter->rank(); + iter = iter->right.get(); + } + else { + break; + } + } + return Iterator{iter}; +} + +// Find iterator such that predicate returns true for all elements from [begin(),iterator) +// If p(*iterator) is true for all elements, this will be end +// Otherwise p(*iterator) will be false and iterator will point to beginning of list or previous element will be true +template template auto RankedTree::find_first_false(const UnaryPredicate& p) const -> ConstIterator { + Node* iter = m_root.get(); + Node* result = nullptr; + while(iter) { + if(p(iter->cdata())) { + iter = iter->right.get(); // Step forward + } + else { + result = iter; // Save the candidate + iter = iter->left.get(); // Step backwards + } + } + return ConstIterator{result}; +} + +template template auto RankedTree::equal_range(const SignedPredicate& p) const -> Range { + Node* iter = m_root.get(); + if(!iter) return {end(), end()}; // Empty range if no root + + const auto scan_left = [&p](Node* iter) { + return iter; + }; + const auto scan_right = [&p](Node* iter) { + while(iter->right && p(iter->right->cdata()) == 0) + iter = iter->right; + return iter; + }; + + for(;;) { + const int comp = p(iter->cdata()); + + if(comp > 0) { + // Target range is strictly to the left of iter + if(!iter->left) { + // Nothing to the left so result is empty range at start of list + assert(iter->prev() == nullptr); + return {Iterator{iter},Iterator{iter}}; + } + iter = iter->left.get(); + } + else if(comp < 0) { + // Target range is strictly to the right of iter + if(!iter->right) { + // Nothing to the right so result is empty range at end of list + assert(iter->next() == nullptr); + return {end(), end()}; + } + iter = iter->right.get(); + } + else { + Node* lo_iter = iter; + Node* hi_iter = iter; + while(lo_iter->left && p(lo_iter->left->cdata()) == 0) + lo_iter = lo_iter->left.get(); + while(hi_iter->right && p(hi_iter->right->cdata()) == 0) + hi_iter = hi_iter->right.get(); + assert(hi_iter->next() == hi_iter->right.get()); // Search started at root so if we stopped at a leaf it must be the max leaf + hi_iter = hi_iter->right.get(); // Bump hi side up by one so to make an open range + return {Iterator{lo_iter}, Iterator{hi_iter}}; + } + } +} + + + +template auto RankedTree::force_red(Node& node) -> void { + if(node.is_red()) { + return; + } + if(node.is_root()) { + // Can safely swap root color + node.set_color(Color::RED); + return; + } + assert(node.is_black()); // Not red so must be black + assert(node.parent->left && node.parent->right); // Black nodes should have a sibling + if(node.parent->left->is_black() && node.parent->right->is_black()) { + force_red(*node.parent); + flip_color(*node.parent); + return; + } + // Sibling must be red (since this node is black and nodes aren't both black) + // From left leaning property this node must also be on the right + assert(node.is_right()); + assert(node.parent->left->is_red()); + rotate_right(parent_link(*node.parent)); + // We just broke left leaning property on parent, but now we fix it + assert(node.parent->is_red()); + assert(node.parent->left && node.parent->left->is_black()); + assert(node.parent->right.get() == &node); + assert(node.is_black()); // Still + flip_color(*node.parent); // Now we can make this node red + assert(node.is_red()); +} + +template auto RankedTree::extract_node(Node* node) -> void { + assert(node); + assert(!node->right || (node->right && node->left)); // Check that node is left-leaning since we're about to use that fact + + if(node->left) { + // Node has children. Need to shuffle things around so that we are erasing a node without children + Node& alternate = node->left->min_child(); // Find some node below h without children + assert(!alternate.left && !alternate.right); // Since tree is left-leaning, min_child should have no children + // Now we swap position of alternate with position of h + swap_positions(alternate, *node); + } + + // Step above ensures we are only trying to erase a node without children + assert(!node->left && !node->right); + + force_red(*node); + + // Update rank for all ancestors + for(Node* iter = node; iter->parent; iter = iter->parent) { + if(iter->is_left()) iter->parent->update_rank(-1); + } + + if(node->is_root()) { + assert(m_root.get() == node); + // Delete the node by resetting the root + m_root.reset(); + // No cleanup needed since tree is now empty + } + else { + // Grab link to the node that we're about to destroy + NodeLink& link_to_node = parent_link(*node); + // Grab link above that which we can use for cleanup + NodeLink& link_to_parent = parent_link(*node->parent); + // Delete the node by resetting the link to it + link_to_node.reset(); + + // Cleanup as needed back to the root + fixup_ancestors(link_to_parent); + } +} + +} // geode namespace \ No newline at end of file diff --git a/geode/structure/module.cpp b/geode/structure/module.cpp index e014da18..bf936c47 100644 --- a/geode/structure/module.cpp +++ b/geode/structure/module.cpp @@ -3,4 +3,5 @@ using namespace geode; void wrap_structure() { GEODE_WRAP(heap) + GEODE_WRAP(ranked_tree) } diff --git a/geode/structure/test_ranked_tree.py b/geode/structure/test_ranked_tree.py new file mode 100644 index 00000000..4f50e188 --- /dev/null +++ b/geode/structure/test_ranked_tree.py @@ -0,0 +1,8 @@ +from geode import * +from numpy import random + +def test_ranked_tree(): + ranked_tree_test() + +if __name__=='__main__': + test_ranked_tree() From 2eb3aaeeb3b47d8fcae57a4dad099b0935df480b Mon Sep 17 00:00:00 2001 From: Ruza Markov Date: Fri, 7 Apr 2017 15:02:58 -0700 Subject: [PATCH 2/2] typo --- geode/structure/RankedTree_p.h | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/geode/structure/RankedTree_p.h b/geode/structure/RankedTree_p.h index 7798ae97..442b2056 100644 --- a/geode/structure/RankedTree_p.h +++ b/geode/structure/RankedTree_p.h @@ -95,7 +95,7 @@ template auto RankedTree::test_subtree_invarients(const Node* node) } result.max_height = 1 + max(l_stats.max_height, r_stats.max_height); GEODE_ASSERT(2*result.black_height >= result.max_height); // Should have at least every other link black - GEODE_ASSERT(result.total_size <= 1< auto RankedTree::extract_node(Node* node) -> void { } } -} // geode namespace \ No newline at end of file +} // geode namespace