Update Tree structures in compactds

mourisl · Jul 11, 2024 · d6b9428 · d6b9428
1 parent 1373c0e
commit d6b9428
Showing 1 changed file with 25 additions and 13 deletions.
diff --git a/compactds/Tree_Plain.hpp b/compactds/Tree_Plain.hpp
@@ -46,15 +46,25 @@ class Tree_Plain: public Tree
 {
 private:
   std::vector<struct _plainTreeNode> _nodes ;
+  size_t _root ; // allow flexible root setting. For compact representation, this 
 public:
-  Tree_Plain() {}
+  Tree_Plain() 
+  {
+    _root = 0 ;
+  }
+
   ~Tree_Plain() {}
 
+  void SetRoot(size_t r)
+  {
+    _root = r ;
+  }
+
   void Init()
   {
     _n = 1 ;
 
-    struct _plainTreeNode node(0, 0, 0, 0) ;
+    struct _plainTreeNode node(_root, _root, _root, _root) ;
     _nodes.push_back(node) ;
   }
 
@@ -66,7 +76,7 @@ class Tree_Plain: public Tree
     _n = n ;
     for (i = 0 ; i < _n ; ++i)
     {
-      struct _plainTreeNode node(0, 0, 0, 0) ;
+      struct _plainTreeNode node(_root, _root, _root, _root) ;
       _nodes.push_back(node) ;
     }
   }
@@ -81,10 +91,10 @@ class Tree_Plain: public Tree
   size_t AddNode(size_t parent)
   {
     size_t id = _nodes.size() ; 
-    struct _plainTreeNode node(parent, 0, 0, 0) ;
+    struct _plainTreeNode node(parent, _root, _root, _root) ;
     size_t lastSibling = LastChild(parent) ;
 
-    if (lastSibling == 0)
+    if (lastSibling == _root)
       _nodes[parent].child = id ;
     else
       _nodes[lastSibling].sibling = id ;
@@ -101,7 +111,7 @@ class Tree_Plain: public Tree
     _nodes[c].parent = parent ;
 
     size_t lastSibling = LastChild(parent) ;
-    if (lastSibling == 0)
+    if (lastSibling == _root)
       _nodes[parent].child = c ;
     else
       _nodes[lastSibling].sibling = c ;
@@ -111,7 +121,7 @@ class Tree_Plain: public Tree
 
   size_t Root() const 
   {
-    return 0 ;
+    return _root ;
   }
 
   // t-th(1-based) child ;
@@ -140,7 +150,7 @@ class Tree_Plain: public Tree
   {
     size_t i ;
     size_t c = FirstChild(v) ;
-    for (i = 0 ; c != 0 ; ++i)
+    for (i = 0 ; c != _root ; ++i)
       c = NextSibling(c) ;
     return i ;
   }
@@ -149,7 +159,7 @@ class Tree_Plain: public Tree
   {
     std::vector<size_t> ret ;
     size_t c = _nodes[v].child ;
-    while (c != 0)
+    while (c != _root)
     {
       ret.push_back(c) ;
       c = _nodes[c].sibling ;
@@ -190,7 +200,7 @@ class Tree_Plain: public Tree
 
   bool IsLeaf(size_t v) const
   {
-    if (_nodes[v].child == 0)
+    if (_nodes[v].child == _root)
       return true ;
     return false ;
   }
@@ -219,8 +229,8 @@ class Tree_Plain: public Tree
   size_t ChildrenLabeled(size_t v, size_t l) const  
   {
     size_t ret = 0 ;
-    int c = _nodes[v].child ;
-    while (c != 0)
+    size_t c = _nodes[v].child ;
+    while (c != _root)
     {
       if (_nodes[c].label == l)
         ++ret ;
@@ -234,7 +244,7 @@ class Tree_Plain: public Tree
   {
     size_t cnt = 0 ;
     size_t c = _nodes[v].child ;
-    while (c != 0)
+    while (c != _root)
     {
       if (_nodes[c].label == l)
         ++cnt ;
@@ -255,6 +265,7 @@ class Tree_Plain: public Tree
   {
     Tree::Save(fp) ;
     size_t i ;
+    SAVE_VAR(fp, _root) ;
     for (i = 0 ; i < _n ; ++i)
       _nodes[i].Save(fp) ;
   }
@@ -264,6 +275,7 @@ class Tree_Plain: public Tree
     std::vector< struct _plainTreeNode >().swap(_nodes) ;    
 
     Tree::Load(fp) ;
+    LOAD_VAR(fp, _root) ;
     size_t i ;
     struct _plainTreeNode node(0, 0, 0, 0) ;
     for (i = 0 ; i < _n ; ++i)