Create move-sub-tree-of-n-ary-tree.cpp

Author: kamyu104

C++/move-sub-tree-of-n-ary-tree.cpp

@@ -0,0 +1,239 @@
+// Time:  O(n)
+// Space: O(h)
+
+/*
+// Definition for a Node.
+class Node {
+public:
+    int val;
+    vector<Node*> children;
+
+    Node() {}
+
+    Node(int _val) {
+        val = _val;
+    }
+
+    Node(int _val, vector<Node*> _children) {
+        val = _val;
+        children = _children;
+    }
+};
+*/
+
+// one pass solution without recursion
+class Solution {
+public:
+    Node* moveSubTree(Node* root, Node* p, Node* q) {
+        unordered_map<Node *, Node *> lookup;
+        const auto& is_ancestor = iter_find_parents(root, nullptr, p, q, false, &lookup);
+        if (lookup.count(p) && lookup[p] == q) {
+            return root;
+        }
+        q->children.emplace_back(p);
+        if (!is_ancestor) {
+            lookup[p]->children.erase(find(begin(lookup[p]->children), end(lookup[p]->children), p));
+        } else {
+            lookup[q]->children.erase(find(begin(lookup[q]->children), end(lookup[q]->children), q));
+            if (p == root) {
+                root = q;
+            } else {
+                *find(begin(lookup[p]->children), end(lookup[p]->children), p) = q;
+            }
+        }
+        return root;
+    }
+
+private:
+    bool iter_find_parents(Node *node, Node *parent, Node *p, Node *q,
+                           bool is_ancestor,
+                           unordered_map<Node *, Node *> *lookup) {
+        vector<tuple<int, Node *, Node *, bool, int>> stk = {tuple(1, node, parent, is_ancestor, -1)};
+        while (!stk.empty()) {
+            const auto [step, node, parent, is_ancestor, i] = stk.back(); stk.pop_back();
+            if (step == 1) {
+                if (node == p || node == q) {
+                    (*lookup)[node] = parent;
+                    if (lookup->size() == 2) {
+                        return is_ancestor;
+                    }
+                }
+                stk.emplace_back(2, node, parent, is_ancestor, node->children.size() - 1);
+            } else {
+                if (i < 0) {
+                    continue;
+                }
+                stk.emplace_back(2, node, parent, is_ancestor, i - 1);
+                stk.emplace_back(1, node->children[i], node, is_ancestor || node == p, -1);
+            }
+        }
+        assert(false);
+        return false;
+    }
+};
+
+// Time:  O(n)
+// Space: O(h)
+// one pass solution with recursion
+class Solution_Recu {
+public:
+    Node* moveSubTree(Node* root, Node* p, Node* q) {
+        unordered_map<Node *, Node *> lookup;
+        const auto& [_, is_ancestor] = find_parents(root, nullptr, p, q, false, &lookup);
+        if (lookup.count(p) && lookup[p] == q) {
+            return root;
+        }
+        q->children.emplace_back(p);
+        if (!is_ancestor) {
+            lookup[p]->children.erase(find(begin(lookup[p]->children), end(lookup[p]->children), p));
+        } else {
+            lookup[q]->children.erase(find(begin(lookup[q]->children), end(lookup[q]->children), q));
+            if (p == root) {
+                root = q;
+            } else {
+                *find(begin(lookup[p]->children), end(lookup[p]->children), p) = q;
+            }
+        }
+        return root;
+    }
+
+private:
+    pair<bool, bool> find_parents(Node *node, Node *parent, Node *p, Node *q,
+                                  bool is_ancestor,
+                                  unordered_map<Node *, Node *> *lookup) {
+        if (node == p || node == q) {
+            (*lookup)[node] = parent;
+            if (lookup->size() == 2) {
+                return {true, is_ancestor};
+            }
+        }
+        for (const auto& child : node->children) {
+            const auto& [found, result] = find_parents(child, node, p, q, is_ancestor || node == p, lookup);
+            if (found) {
+                return {true, result};
+            }
+        }
+        return {false, false};
+    }
+};
+
+// Time:  O(n)
+// Space: O(h)
+// two pass solution without recursion
+class Solution2 {
+public:
+    Node* moveSubTree(Node* root, Node* p, Node* q) {
+        unordered_map<Node *, Node *> lookup;
+        iter_find_parents(root, nullptr, p, q, &lookup);
+        if (lookup.count(p) && lookup[p] == q) {
+            return root;
+        }
+        q->children.emplace_back(p);
+        if (!iter_is_ancestor(p, q)) {
+            lookup[p]->children.erase(find(begin(lookup[p]->children), end(lookup[p]->children), p));
+        } else {
+            lookup[q]->children.erase(find(begin(lookup[q]->children), end(lookup[q]->children), q));
+            if (p == root) {
+                root = q;
+            } else {
+                *find(begin(lookup[p]->children), end(lookup[p]->children), p) = q;
+            }
+        }
+        return root;
+    }
+
+private:
+    void iter_find_parents(Node *node, Node *parent, Node *p, Node *q,
+                           unordered_map<Node *, Node *> *lookup) {
+        vector<tuple<int, Node *, Node *, int>> stk = {tuple(1, node, parent, -1)};
+        while (!stk.empty()) {
+            const auto [step, node, parent, i] = stk.back(); stk.pop_back();
+            if (step == 1) {
+                if (node == p || node == q) {
+                    (*lookup)[node] = parent;
+                    if (lookup->size() == 2) {
+                        return;
+                    }
+                }
+                stk.emplace_back(2, node, parent, node->children.size() - 1);
+            } else {
+                if (i < 0) {
+                    continue;
+                }
+                stk.emplace_back(2, node, parent, i - 1);
+                stk.emplace_back(1, node->children[i], node, -1);
+            }
+        }
+    }
+    
+    bool iter_is_ancestor(Node *node, Node *q) {
+        vector<tuple<int, Node *, int>> stk = {tuple(1, node, -1)};
+        while (!stk.empty()) {
+            const auto [step, node, i] = stk.back(); stk.pop_back();
+            if (step == 1) {
+                stk.emplace_back(2, node, node->children.size() - 1);
+            } else {
+                if (i < 0) {
+                    continue;
+                }
+                if (node->children[i] == q) {
+                    return true;
+                }
+                stk.emplace_back(2, node, i - 1);
+                stk.emplace_back(1, node->children[i], -1);
+            }
+        }
+        return false;
+    }
+};
+
+// Time:  O(n)
+// Space: O(h)
+// two pass solution with recursion
+class Solution2_Recu {
+public:
+    Node* moveSubTree(Node* root, Node* p, Node* q) {
+        unordered_map<Node *, Node *> lookup;
+        find_parents(root, nullptr, p, q, &lookup);
+        if (lookup.count(p) && lookup[p] == q) {
+            return root;
+        }
+        q->children.emplace_back(p);
+        if (!is_ancestor(p, q)) {
+            lookup[p]->children.erase(find(begin(lookup[p]->children), end(lookup[p]->children), p));
+        } else {
+            lookup[q]->children.erase(find(begin(lookup[q]->children), end(lookup[q]->children), q));
+            if (p == root) {
+                root = q;
+            } else {
+                *find(begin(lookup[p]->children), end(lookup[p]->children), p) = q;
+            }
+        }
+        return root;
+    }
+
+private:
+    bool find_parents(Node *node, Node *parent, Node *p, Node *q, unordered_map<Node *, Node *> *lookup) {
+        if (node == p || node == q) {
+            (*lookup)[node] = parent;
+            if (lookup->size() == 2) {
+                return true;
+            }
+        }
+        for (const auto& child : node->children) {
+            if (find_parents(child, node, p, q, lookup)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+    bool is_ancestor(Node *node, Node *q) {
+        for (const auto& child : node->children) {
+            if (child == q || is_ancestor(child, q)) {
+                return true;
+            }
+        }
+        return false;
+    }
+};