Encode N-ary Tree to Binary Tree

Problem Id: 431 Difficulty: Hard Tag: Tree


Intuition

The algorithm I use is that, the left points to a linked list to represent the children list. And the right is been used as the linked list's next. For example, if we have a N-ary tree like below:

         1
     /   \    \
   3      2     4
 /  \
5    6

It could be changed to a binary tree like below:

         1
     /
   3  -   2  -  4
 /
5  - 6

Solution


"""
# Definition for a Node.
class Node:
    def __init__(self, val=None, children=None):
        self.val = val
        self.children = children
"""
"""
# Definition for a binary tree node.
class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None
"""
class Codec:

    # Encodes an n-ary tree to a binary tree.
    def encode(self, node: 'Node') -> TreeNode:
        if node is None:
            return None
        binode = TreeNode(node.val)
        if node.children:
            binode.left = self.encode(node.children[0])
            tmp = binode.left
            for child in node.children[1:]:
                tmp.right = self.encode(child)
                tmp = tmp.right
        return binode

    # Decodes your binary tree to an n-ary tree.
    def decode(self, binode: TreeNode) -> 'Node':
        if binode is None:
            return None
        node = Node(binode.val, [])
        tmp = binode.left
        while tmp:
            node.children.append(self.decode(tmp))
            tmp = tmp.right
        return node

# Your Codec object will be instantiated and called as such:
# codec = Codec()
# codec.decode(codec.encode(root))