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

segment_tree.h

@@ -15,7 +15,6 @@ struct TreeNode {
     Coordinate right_coord;
     Unit covered_length;
     inline Unit segment_length();
-    inline bool covered();
 };
 
 
@@ -31,9 +30,8 @@ public:
     SegmentTree(const std::vector<RectCoord>& coords);
 
     inline Unit length_covered_intervals();
-    void add_interval(Interval interval);
-    void remove_interval(Interval interval);
-    void check_integrity();
+    inline void add_interval(Interval interval);
+    inline void remove_interval(Interval interval);
 
 private:
     inline static Index left_child_idx(Index node_idx);
@@ -43,6 +41,8 @@ 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,16 +65,22 @@ 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);
-    _nodes[node_idx].covered_length = _nodes[node_idx].segment_length();
+    update_covered_length(node_idx);
 }
 
 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) {
@@ -86,8 +92,9 @@ Index SegmentTree::right_child_idx(Index node_idx) {
 }
 
 void SegmentTree::update_covered_length(Index node_idx) {
-    assert(!_nodes[node_idx].covered());
-    if (is_leaf(node_idx)) {
+    if (_nodes[node_idx].coverage > 0) {
+        _nodes[node_idx].covered_length = _nodes[node_idx].segment_length();
+    } else if (is_leaf(node_idx)) {
         _nodes[node_idx].covered_length = 0;
     } else {
         _nodes[node_idx].covered_length = \
@@ -99,14 +106,9 @@ 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