Skip to content

1516. Move Sub-Tree of N-Ary Tree 👎

  • Time: $O(n)$
  • Space: $O(h)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
class Solution {
 public:
  Node* moveSubTree(Node* root, Node* p, Node* q) {
    if (find(q->children.begin(), q->children.end(), p) != q->children.end())
      return root;

    // Create a dummy node for the case when root == p.
    Node* dummy = new Node(0, {root});

    // Get each parent of p and q.
    Node* pParent = getParent(dummy, p);
    Node* qParent = getParent(p, q);

    // Get p's original index in p's parent.
    vector<Node*>& pSiblings = pParent->children;
    const int pIndex =
        find(pSiblings.begin(), pSiblings.end(), p) - pSiblings.begin();
    pSiblings.erase(pSiblings.begin() + pIndex);

    q->children.push_back(p);

    // If q is in p's subtree, qParent != nullptr.
    if (qParent != nullptr) {
      vector<Node*>& qSiblings = qParent->children;
      std::erase(qSiblings, q);
      pSiblings.insert(pSiblings.begin() + pIndex, q);
    }

    return dummy->children[0];
  }

 private:
  Node* getParent(Node* root, Node* target) {
    for (Node* child : root->children) {
      if (child == target)
        return root;
      Node* res = getParent(child, target);
      if (res != nullptr)
        return res;
    }
    return nullptr;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
class Solution:
  def moveSubTree(self, root: 'Node', p: 'Node', q: 'Node') -> 'Node':
    if p in q.children:
      return root

    # Create a dummy Node for the case when root == p
    dummy = Node(None, [root])

    # Get each parent of p and q
    pParent = self._getParent(dummy, p)
    qParent = self._getParent(p, q)

    # Get p's original index in p's parent
    pIndex = pParent.children.index(p)
    pParent.children.pop(pIndex)

    q.children.append(p)

    # If q is in the p's subtree, qParent != None
    if qParent:
      qParent.children.remove(q)
      pParent.children.insert(pIndex, q)

    return dummy.children[0]

  def _getParent(self, root: 'Node', target: 'Node') -> Optional['Node']:
    for child in root.children:
      if child == target:
        return root
      res = self._getParent(child, target)
      if res:
        return res
    return None