From a0f3b12ecec2de74623d485c33166f7027fd79d2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maximilian=20Ke=C3=9Fler?= <git@maximilian-kessler.de>
Date: Sun, 17 Apr 2022 13:38:52 +0200
Subject: [PATCH] implement segment tree with stack directly

---
 segment_tree.cpp | 65 +++++++++++++++++++++++++++++++++---------------
 segment_tree.h   | 25 +++++++------------
 2 files changed, 54 insertions(+), 36 deletions(-)

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 <cassert>
 #include <bit>
+#include <stack>
 #include "segment_tree.h"
 
 SegmentTree::SegmentTree(const std::vector<RectCoord> &coords):
@@ -33,30 +34,54 @@ 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) {
-    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<Index> visit_to;
+    std::vector<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) {
+                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<Index> visit_to;
+    std::vector<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) {
+                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<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);
 
 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;
 }