Problem statement

https://leetcode.com/problems/move-sub-tree-of-n-ary-tree/

Solution

Not very interesting problem, where we need to carefully check all possible cases and handle them. Let dfs(node, target) be function which traverse subtree of node and look for target node in this subtree: it return the parent of node target. (not mine code, I think it is better tu use dfs with parent variable)

Evaluate q_parent = dfs(p, q) and p_parent = dfs(root, p). Then we can have different options:

  1. If q_parent != None and p_parent != None, it means, that q is children of p and p is not equal to root (Example 4 in problem description). We need to remove q from children of q_parent and append p to children of q. Then we reattach subtree q to the place we deleted it from.
  2. If q_parent = None and p_parent == None, it means that q is children of p and p is root of our tree (Example 5 in problem descripton). We need to remove q from children of q_parent and append p to children of q. We do not need to do the second part like in previous case, but change root to q.
  3. If q_parent == None and p_parent != q, it means that q is not chilren of p and than p is not direct child of q (Examples 1 and 3 in problem description). In this case we just remove p from childrens of p_parent and add p to childrens of q. Note that in example 1 node p is in subtree of q and in example 3 they have LSA root, but in fact these two examples are quite the same.

Complexity

It is O(n) for time, where n is size of tree and O(h) for space.

Code

class Solution:
    def moveSubTree(self, root, p, q):
        def dfs(node, target):
            if not node: return node
            for child in node.children:
                if child == target: return node
                p = dfs(child, target)        
                if p: return p
            return None

        q_parent = dfs(p, q)
        p_parent = dfs(root, p)

        if q_parent:
            q_parent.children.remove(q)
            q.children.append(p)
            if p_parent:
                j = p_parent.children.index(p)
                p_parent.children[j] = q
            else:
                root = q
        elif p_parent != q:
            p_parent.children.remove(p)
            q.children.append(p)

        return root