Leetcode 0431 Encode N-ary Tree to Binary Tree

Problem statement

https://leetcode.com/problems/encode-n-ary-tree-to-binary-tree/

Solution

Very difficult problem, if you see it first time. The strategy follows two rules:

1. The left child of each node in the binary tree is the next sibling of the node in the N-ary tree.
2. The right child of each node in the binary tree is the first child of the node in the N-ary tree.

Also, we can easily recover the N-ary tree from the binary tree we converted. The recursion recovery strategy for each node is:

1. Deal with its children recursively.
2. Add its left child as the next child of its parent if it has a left child.
3. Add its right child as the first child of the node itself if it has a right child.

Also see Wikipedia for more visualization.

Complexity

Time complexity is O(n), where n is number of nodes.

Code

class Codec:
    def encode(self, root):
        if not root: return None
        res = TreeNode(root.val)
        
        if root.children:
            res.left = self.encode(root.children[0])
            
        cur = res.left
        
        for i in range(1, len(root.children)):
            cur.right = self.encode(root.children[i])
            cur = cur.right
            
        return res

    def decode(self, root):
        if not root: return None
        res = Node(root.val, [])
        cur = root.left
        while cur:
            res.children.append(self.decode(cur))
            cur = cur.right

        return res