only update removed node if not covered anymore

main.cpp

@@ -22,10 +22,10 @@ int main() {
     srand(987192345);
 
     add_rect(0,0,5,4);
-    check_area(20);
+  //  check_area(20);
 
     add_rect(3,2,6,6);
-    check_area(28);
+//    check_area(28);
 
     add_rect(1,-2,8,4);
     check_area(52);

segment_tree.cpp

@@ -3,6 +3,7 @@
 //
 #include <cassert>
 #include <bit>
+#include <stack>
 #include "segment_tree.h"
 
 SegmentTree::SegmentTree(const std::vector<RectCoord> &coords):
@@ -33,30 +34,75 @@ SegmentTree::SegmentTree(const std::vector<RectCoord> &coords):
     assert(_nodes.front().segment_length() == coords.back().coord - coords.front().coord);
 }
 
-void SegmentTree::add_interval(Interval interval, Index node_idx) {
+void SegmentTree::add_interval(Interval interval) {
+    check_integrity();
+    std::stack<Index> visit_to;
+    std::stack<Index> 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) {
-            add_interval(interval, left_child_idx(node_idx));
+                visit_to.push(left_child_idx(node_idx));
             }
             if(right_child(node_idx).left_coord < interval.right) {
-            add_interval(interval, right_child_idx(node_idx));
+                visit_to.push(right_child_idx(node_idx));
             }
-        update_covered_length(node_idx);
+            if(!_nodes[node_idx].covered()) {
+                update_to.push(node_idx);
             }
         }
+    }
+    while(!update_to.empty()) {
+        update_covered_length(update_to.top());
+        update_to.pop();
+    }
+    check_integrity();
+}
 
-void SegmentTree::remove_interval(Interval interval, Index node_idx) {
+void SegmentTree::remove_interval(Interval interval) {
+    check_integrity();
+    std::stack<Index> visit_to;
+    std::stack<Index> 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) {
-            remove_interval(interval, left_child_idx(node_idx));
+                visit_to.push(left_child_idx(node_idx));
             }
             if(right_child(node_idx).left_coord < interval.right) {
-            remove_interval(interval, right_child_idx(node_idx));
+                visit_to.push(right_child_idx(node_idx));
+            }
+            if(!_nodes[node_idx].covered()) {
+                update_to.push(node_idx);
+            }
+        }
+    }
+    while(!update_to.empty()) {
+        update_covered_length(update_to.top());
+        update_to.pop();
+    }
+    check_integrity();
+}
+
+void SegmentTree::check_integrity() {
+    for(Index node_idx = 0 ; node_idx < _nodes.size(); ++node_idx) {
+        if(_nodes[node_idx].covered()) {
+            assert(_nodes[node_idx].covered_length == _nodes[node_idx].segment_length());
+        } else {
+            if(is_leaf(node_idx)) {
+                assert(_nodes[node_idx].covered_length == 0);
+            } else {
+                assert(_nodes[node_idx].covered_length == \
+                left_child(node_idx).covered_length + right_child(node_idx).covered_length);
+            }
         }
-        update_covered_length(node_idx);
     }
 }

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,9 @@ public:
     SegmentTree(const std::vector<RectCoord>& 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);
+    void check_integrity();
 
 private:
     inline static Index left_child_idx(Index node_idx);
@@ -41,8 +43,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,23 +65,17 @@ 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) {
     --(_nodes[node_idx].coverage);
+    if(!_nodes[node_idx].covered()) {
         update_covered_length(node_idx);
     }
+}
 
 Index SegmentTree::left_child_idx(Index node_idx) {
     return 2*node_idx + 1;
@@ -92,9 +86,8 @@ 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)) {
+    assert(!_nodes[node_idx].covered());
+    if (is_leaf(node_idx)) {
         _nodes[node_idx].covered_length = 0;
     } else {
         _nodes[node_idx].covered_length = \
@@ -106,9 +99,14 @@ 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;
+    return node_idx >= _num_meta_nodes;
 }
 
 
+
 #endif //PROG_SEGMENT_TREE_H