diff --git a/segment_tree.cpp b/segment_tree.cpp index 86f7c7f..6fcd570 100644 --- a/segment_tree.cpp +++ b/segment_tree.cpp @@ -3,6 +3,7 @@ // #include #include +#include #include "segment_tree.h" SegmentTree::SegmentTree(const std::vector &coords): @@ -33,30 +34,54 @@ SegmentTree::SegmentTree(const std::vector &coords): assert(_nodes.front().segment_length() == coords.back().coord - coords.front().coord); } -void SegmentTree::add_interval(Interval interval, Index node_idx) { - if (interval.left <= _nodes[node_idx].left_coord && _nodes[node_idx].right_coord <= interval.right) { - add_coverage(node_idx); - } else { - if(interval.left < left_child(node_idx).right_coord) { - add_interval(interval, left_child_idx(node_idx)); +void SegmentTree::add_interval(Interval interval) { + std::stack visit_to; + std::vector update_to; + visit_to.push(0); + while(!visit_to.empty()) { + Index node_idx = visit_to.top(); + visit_to.pop(); + if (interval.left <= _nodes[node_idx].left_coord && _nodes[node_idx].right_coord <= interval.right) { + add_coverage(node_idx); + } else { + if(interval.left < left_child(node_idx).right_coord) { + visit_to.push(left_child_idx(node_idx)); + } + if(right_child(node_idx).left_coord < interval.right) { + visit_to.push(right_child_idx(node_idx)); + } + if(!_nodes[node_idx].covered()) { + update_to.push_back(node_idx); + } } - if(right_child(node_idx).left_coord < interval.right) { - add_interval(interval, right_child_idx(node_idx)); - } - update_covered_length(node_idx); + } + for(auto update_idx : update_to) { + update_covered_length(update_idx); } } -void SegmentTree::remove_interval(Interval interval, Index node_idx) { - if (interval.left <= _nodes[node_idx].left_coord && _nodes[node_idx].right_coord <= interval.right) { - remove_coverage(node_idx); - } else { - if(interval.left < left_child(node_idx).right_coord) { - remove_interval(interval, left_child_idx(node_idx)); +void SegmentTree::remove_interval(Interval interval) { + std::stack visit_to; + std::vector update_to; + visit_to.push(0); + while(!visit_to.empty()) { + Index node_idx = visit_to.top(); + visit_to.pop(); + if (interval.left <= _nodes[node_idx].left_coord && _nodes[node_idx].right_coord <= interval.right) { + remove_coverage(node_idx); + } else { + if(interval.left < left_child(node_idx).right_coord) { + visit_to.push(left_child_idx(node_idx)); + } + if(right_child(node_idx).left_coord < interval.right) { + visit_to.push(right_child_idx(node_idx)); + } + if(!_nodes[node_idx].covered()) { + update_to.push_back(node_idx); + } } - if(right_child(node_idx).left_coord < interval.right) { - remove_interval(interval, right_child_idx(node_idx)); - } - update_covered_length(node_idx); + } + for(auto update_idx : update_to) { + update_covered_length(update_idx); } } \ No newline at end of file diff --git a/segment_tree.h b/segment_tree.h index 6b929df..9abb318 100644 --- a/segment_tree.h +++ b/segment_tree.h @@ -15,6 +15,7 @@ struct TreeNode { Coordinate right_coord; Unit covered_length; inline Unit segment_length(); + inline bool covered(); }; @@ -30,8 +31,8 @@ public: SegmentTree(const std::vector& coords); inline Unit length_covered_intervals(); - inline void add_interval(Interval interval); - inline void remove_interval(Interval interval); + void add_interval(Interval interval); + void remove_interval(Interval interval); private: inline static Index left_child_idx(Index node_idx); @@ -41,8 +42,6 @@ private: inline void add_coverage(Index node_idx); inline void remove_coverage(Index node_idx); inline void update_covered_length(Index node_idx); - void add_interval(Interval interval, Index node_idx); - void remove_interval(Interval interval, Index node_idx); inline bool is_leaf(Index node_idx); }; @@ -65,17 +64,9 @@ Unit SegmentTree::length_covered_intervals() { return _nodes.front().covered_length; } -void SegmentTree::add_interval(Interval interval) { - add_interval(interval, 0); -} - -void SegmentTree::remove_interval(Interval interval) { - remove_interval(interval, 0); -} - void SegmentTree::add_coverage(Index node_idx) { ++(_nodes[node_idx].coverage); - update_covered_length(node_idx); + _nodes[node_idx].covered_length = _nodes[node_idx].segment_length(); } void SegmentTree::remove_coverage(Index node_idx) { @@ -92,9 +83,7 @@ Index SegmentTree::right_child_idx(Index node_idx) { } void SegmentTree::update_covered_length(Index node_idx) { - if (_nodes[node_idx].coverage > 0) { - _nodes[node_idx].covered_length = _nodes[node_idx].segment_length(); - } else if (is_leaf(node_idx)) { + if (is_leaf(node_idx)) { _nodes[node_idx].covered_length = 0; } else { _nodes[node_idx].covered_length = \ @@ -106,6 +95,10 @@ Unit TreeNode::segment_length() { return right_coord - left_coord; } +bool TreeNode::covered() { + return coverage > 0; +} + bool SegmentTree::is_leaf(Index node_idx) { return node_idx > _num_meta_nodes; }