1516. Move Sub-Tree of N-Ary Tree

Problem Explanation

This problem involves dealing with an N-ary tree where each node can have N children. We have the root of the tree and two nodes, p and q. The task is to move the subtree of node p, under node q. Node p should become the last child of node q.

There are three scenarios to consider:

1. Node q is in the subtree of node p.
2. Node p is in the subtree of node q.
3. Neither node p nor q lie in each other's subtree.

For the second and third cases, the problem is straightforward - move p (with its subtree) to be a child of q. But in the first case, doing so may disconnect the tree. Therefore, we must reconnect the tree again. In this case, we first move p with all of its subtree except q and add it as a child to q. Then we see if the tree is disconnected and reconnect it accordingly.

Walkthrough

Let's take an example where our input tree is [1,null,2,3,null,4,5,null,6,null,7,8], p equals 2 and q equals 7. This example falls under the first case where q is in the subtree of p.

First step: We move p (with all of its subtree except q) and add it as a child to q. This gives us a tree [1,null,7,3,null,2,null,6,null,4,5,null,null,8].

Then we notice that the tree is disconnected so we need to reconnect q to replace p.

Approach

1. The first step in the solution is checking if p already is a child of q. In this case, we don't have to do anything.

2. We add a dummy node to handle the case where p is the root of the tree.

3. We get the parent of p and remove it from its children list.

4. Add p as a child to q.

5. We check if q is in the subtree of p. If it is, we have to also update q's parent to replace q with p.

Python Solution

1
2python
3class Solution:
4    def moveSubTree(self, root, p, q):
5        if p in q.children:
6            return root
7        
8        dummy = Node(0, [root])
9        
10        pParent = self.getParent(dummy, p)
11        pParent.children.remove(p)
12        
13        q.children.append(p)
14        
15        qParent = self.getParent(p, q)
16        if qParent:
17            qParent.children.remove(q)
18            pParent.children.append(q)
19            
20        return dummy.children[0]
21          
22    def getParent(self, root, node):
23        stack = [(root, None)]
24        while stack:
25            curr, parent = stack.pop()
26            if curr == node:
27                return parent
28            for child in curr.children:
29                stack.append((child, curr))
30        return None

Java Solution

1
2java
3public class Solution {
4    public Node moveSubTree(Node root, Node p, Node q) {
5        if (q.children.contains(p)) {
6            return root;
7        }
8
9        Node dummy = new Node(0, Arrays.asList(root));
10
11        Node pParent = getParent(dummy, p);
12        pParent.children.remove(p);
13        
14        q.children.add(p);
15        
16        Node qParent = getParent(p, q);
17        if (qParent != null) {
18            qParent.children.remove(q);
19            pParent.children.add(q);
20        }
21
22        return dummy.children.get(0);
23    }
24    private Node getParent(Node root, Node target) {
25        Deque<Pair<Node, Node>> stack = new ArrayDeque<>();
26        stack.push(new Pair(root, null));
27        while (!stack.isEmpty()) {
28            Pair<Node, Node> curr = stack.pop();
29            if (curr.getKey() == target) {
30                return curr.getValue();
31            }
32            for (Node child: curr.getKey().children) {
33                stack.push(new Pair(child, curr.getKey()));
34            }
35        }
36        return null;
37    }
38}

C# Solution

1
2csharp
3public class Solution {
4    public Node MoveSubTree(Node root, Node p, Node q) {
5        if (q.children.Contains(p)) {
6            return root;
7        }
8        
9        Node dummy = new Node(0, new List<Node> {root});
10        
11        Node pParent = GetParent(dummy, p);
12        pParent.children.Remove(p);
13        
14        q.children.Add(p);
15        
16        Node qParent = GetParent(p, q);
17        if (qParent != null) {
18            qParent.children.Remove(q);
19            pParent.children.Add(q);
20        }
21        
22        return dummy.children[0];
23    }
24    
25    private Node GetParent(Node root, Node target) {
26        Stack<Tuple<Node, Node>> stack = new Stack<Tuple<Node, Node>>();
27        stack.Push(new Tuple<Node, Node>(root, null));
28        while (stack.Count > 0) {
29            Tuple<Node, Node> curr = stack.Pop();
30            if (curr.Item1 == target) {
31                return curr.Item2;
32            }
33            foreach (Node child in curr.Item1.children) {
34                stack.Push(new Tuple<Node, Node>(child, curr.Item1));
35            }
36        }
37        return null;
38    }
39}

JavaScript Solution

1
2javascript
3class Solution {
4    moveSubTree(root, p, q){
5        if(q.children.includes(p)){
6            return root;
7        }
8        let dummy = {val: 0, children: [root]};
9        let pParent = this.getParent(dummy, p);
10        pParent.children = pParent.children.filter(child => child !== p);
11        q.children.push(p);
12        let qParent = this.getParent(p, q);
13        if(qParent){
14            qParent.children = qParent.children.filter(child => child !== q);
15            pParent.children.push(q);
16        }
17        return dummy.children[0];
18    }
19    getParent(root, target){
20        let stack = [[root, null]];
21        while(stack.length){
22            let [node, parent] = stack.pop();
23            if(node === target){
24                return parent;
25            }
26            for(let child of node.children){
27                stack.push([child, node]);
28            }
29        }
30        return null;
31    }
32}

C++ Solution

1
2cpp
3class Solution{
4public:
5    Node* moveSubTree(Node* root, Node* p, Node* q){
6        if(find(q->children.begin(), q->children.end(), p) != q->children.end()){
7            return root;
8        }
9        Node* dummy = new Node(0, {root});
10        Node* pParent = getParent(dummy, p);
11        pParent->children.erase(remove(pParent->children.begin(), pParent->children.end(), p), pParent->children.end());
12        q->children.push_back(p);
13        Node* qParent = getParent(p, q);
14        if(qParent){
15            qParent->children.erase(remove(qParent->children.begin(), qParent->children.end(), q), qParent->children.end());
16            pParent->children.push_back(q);
17        }
18        return dummy->children[0];
19    }
20private:
21    Node* getParent(Node* root, Node* target){
22        stack<pair<Node*, Node*>> stk;
23        stk.push({root, nullptr});
24        while(!stk.empty()){
25            auto [node, parent] = stk.top();
26            stk.pop();
27            if(node == target){
28                return parent;
29            }
30            for(auto& child : node->children){
31                stk.push({child, node});
32            }
33        }
34        return nullptr; 
35    }
36};

In this C++ solution, we start by checking if p is already a child of q. If it is, we simply return the root as no moving is needed.

We then create a dummy node and add root to it. We find the parent of p and remove p from its children. We push p into q's children.

Finally, we find if q is present in the subtree of p. If it is, we find the parent of q and remove q from its children and add q to pParent's children.

This way, we have moved the subtree rooted at p to become the last child of q.

The method getParent finds the parent of the target node by DFS (Depth-First Search) using a stack, which stores pairs of current nodes and their parents. If it finds the target node in the stack, it pops and returns the parent. If it doesn't find the target, it returns nullptr.

It returns the new root.

This solution assumes that nodes of the tree hold unique values.

Solution Implementation