Update Segment Tree

Author: Colin Lin

Algon/dstruct/SegmentTree.h

@@ -1,82 +1,96 @@
 #pragma once
 #include <functional>
-#include <iterator>
 #include <vector>
 #include "../utility/Exception.h"
 #include "../common.h"
-#include "../utility/Util.h"
 namespace colinli {
 namespace algon {
 
-  using std::vector;
+using std::vector;
 
-  template<typename TValue, typename TKey>
-  class SegmentTree
-  {
-  public:
-    typedef std::function<TValue(TKey)> Transformer;
-    typedef std::function<TValue(TValue, TValue)> Reducer;
-    SegmentTree(vector<TKey>& collection, Reducer reducer, Transformer trans) {
-      vector_ = collection;
-      reducer_ = reducer;
-      trans_ = trans;
-      root_ = build_tree(0, collection.size(), reducer, trans);
-    }
-    SegmentTree(vector<TKey>&& collection, Reducer reducer, Transformer trans) {
-      vector_ = collection;
-      reducer_ = reducer;
-      trans_ = trans;
-      root_ = build_tree(0, collection.size(), reducer, trans);
-    }
+#define LCHILD(i) (((i)<<1)+1)
+#define RCHILD(i) (((i)<<1)+2)
 
-    struct Node
-    {
-      Node* left;
-      Node* right;
-      TValue value;
-      explicit Node(TValue v) :value(v), left(nullptr), right(nullptr){}
-      explicit Node() :left(nullptr), right(nullptr), value(TValue()){}
-    };
-  private:
-    Node* root_;
-    vector<TKey> vector_;
-    Transformer trans_;
-    Reducer reducer_;
+template <typename TValue, typename TKey>
+class SegmentTree
+{
+public:
+  typedef std::function<TValue(TKey)> Transformer;
+  typedef std::function<TValue(TValue, TValue)> Reducer;
 
-    Node* build_tree(size_t begin,size_t end) {
-      if (begin == end) {
-        return NULL;
-      }
-      auto mid = begin + (end - begin) / 2;
-      Node* node = new Node(trans_(vector_.at(mid)));
-      node->left = build_tree(begin, mid);
-      node->right = build_tree( std::next(mid), end);
-      if (node->left && node->right) {
-        node->value = reducer_(node->left->value, node->right->value);
-      }
-      return node;
-    }
+  SegmentTree(vector<TKey>& collection, Reducer reducer, Transformer trans, TValue initVal):
+    vector_(collection),
+    N(collection.size()),
+    trans_(trans),
+    reducer_(reducer),
+    initVal_(initVal) {
+    node_values_ = new TValue[(N << 2) + 1];
+    build_tree(0, 0, N-1);
+  }
+
+  void Update(size_t pos, TKey newkey) {
+    update(0, 0, N - 1, pos, newkey);
+  }
+
+  TValue Query(size_t begin, size_t end) {
+    return query(0, 0, N - 1, begin, end);
+  }
+
+  virtual ~SegmentTree() {
+    delete[] node_values_;
+  }
+
+private:
+  vector<TKey> vector_;
+  size_t N;
+  Transformer trans_;
+  Reducer reducer_;
+  TValue initVal_;
+  TValue* node_values_;
 
-    void update(Node* node ,size_t begin, size_t end, size_t pos, TKey newkey) {
-      if (begin == end || node == NULL) {
-        return;
-      }
-      if (++begin == end) {
-        node->value = trans_(newkey);
-      }
-      else {
-        auto mid = begin + (begin - end) / 2;
-        if ( pos <= mid ) {
-          update(node->left, begin, mid, pos, newkey);
-          update(node->right, mid+1 , end, pos, newkey);
-          CHECK_THROW(node->left && node->right, NullArgumentError);
-          node->value = reducer_(node->left->value, node->right->value);
-        }
-      }
+  // build substree from x[begin...end] (inclusive range)
+  void build_tree(size_t node_idx, size_t begin, size_t end) {
+    //leaf node 
+    if (begin == end) {
+      node_values_[node_idx] = trans_(vector_.at(begin));
+      return;
     }
-  };
+    auto mid = begin + (end - begin) / 2;
+    build_tree(LCHILD(node_idx), begin, mid);
+    build_tree(RCHILD(node_idx), mid + 1, end);
+    node_values_[node_idx] = reducer_(node_values_[LCHILD(node_idx)], node_values_[RCHILD(node_idx)]);
+  }
 
+  void update(size_t node_idx, size_t begin, size_t end, size_t pos, TKey newkey) {
+    if (begin == end) {
+      node_values_[node_idx] = trans_(newkey);
+      return;
+    }
+    auto mid = begin + (end - begin) / 2;
+    if (pos <= mid) {
+      update(LCHILD(node_idx), begin, mid, pos, newkey);
+    }
+    else {
+      update(RCHILD(node_idx), mid + 1, end, pos, newkey);
+    }
+    node_values_[node_idx] = reducer_(node_values_[LCHILD(node_idx)], node_values_[RCHILD(node_idx)]);
+  }
 
+  TValue query(size_t node_idx, size_t node_begin, size_t node_end, size_t q_begin, size_t q_end) {
+    if (node_begin >= q_begin && node_end <= q_end) {
+      return node_values_[node_idx];
+    }
+    auto mid = node_begin + (node_end - node_begin) / 2;
+    TValue result = initVal_;
+    if (mid > q_begin && node_begin <= q_end) {
+      result = reducer_(result, query(LCHILD(node_idx), node_begin, mid, q_begin, q_end));
+    }
+    if (node_end >= q_begin && mid + 1 <= q_end) {
+      result = reducer_(result, query(RCHILD(node_idx), mid + 1, node_end, q_begin, q_end));
+    }
+    return result;
+  }
+};
 
 }// end of algon ns
 }// end of colinli ns 

Algon/utility/Util.h

@@ -15,18 +15,6 @@ namespace colinli {
   uint32_t sdigits10(int64_t v);
   int ll2string(char *dst, size_t dstlen, long long svalue);
 
-  
-  /// <summary>
-  /// Identity functor 
-  /// </summary>
-  struct identity {
-    template<typename U>
-    constexpr auto operator()(U&& v) const noexcept
-      -> decltype(std::forward<U>(v))
-    {
-      return std::forward<U>(v);
-    }
-  };
 
 
   /// <summary>

tests/SegmentTreeTest.cpp

@@ -0,0 +1,26 @@
+#include "stdafx.h"
+#include "CppUnitTest.h"
+
+using namespace Microsoft::VisualStudio::CppUnitTestFramework;
+using namespace colinli::algon;
+using namespace std;
+namespace tests
+{
+	TEST_CLASS(SegmentTreeTest)
+	{
+	public:
+		
+		TEST_METHOD(TestQueryTree)
+		{
+      auto min_reducer = [](int x, int y) throw() { return x < y ? x : y; };
+      auto identity = [](int x) {return x; };
+      int initv = INT_MAX;
+      vector<int> x{ 1, 2, 3, 4, 5, 6, 7 };
+      SegmentTree<int, int> tree(x, min_reducer, identity, initv);
+      Assert::AreEqual(3, tree.Query(2, 5));
+      tree.Update(3, -1);
+      Assert::AreEqual(-1, tree.Query(1, 6));
+		}
+
+	};
+}

tests/stdafx.h

@@ -10,7 +10,9 @@
 // Headers for CppUnitTest
 #include "CppUnitTest.h"
 #include <string>
+#include <vector>
 // TODO: reference additional headers your program requires here
 #include "../Algon/dstruct/SkipList.h"
 #include "../Algon/dstruct/DisjointSet.h"
-#include "../Algon/algorithm/Numeric.h"
+#include "../Algon/algorithm/Numeric.h"
+#include "../Algon/dstruct/SegmentTree.h"

tests/tests.vcxproj

@@ -88,6 +88,7 @@
   <ItemGroup>
     <ClCompile Include="DisjointSetTest.cpp" />
     <ClCompile Include="NumericTest.cpp" />
+    <ClCompile Include="SegmentTreeTest.cpp" />
     <ClCompile Include="stdafx.cpp">
       <PrecompiledHeader Condition="'$(Configuration)|$(Platform)'=='Debug|Win32'">Create</PrecompiledHeader>
       <PrecompiledHeader Condition="'$(Configuration)|$(Platform)'=='Release|Win32'">Create</PrecompiledHeader>

tests/tests.vcxproj.filters

@@ -35,5 +35,8 @@
     <ClCompile Include="NumericTest.cpp">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="SegmentTreeTest.cpp">
+      <Filter>Source Files</Filter>
+    </ClCompile>
   </ItemGroup>
 </Project>