判断两颗多叉树关于叶子节点的前序遍历是否相等

multi_postorder.py

class TreeNode:
    def __init__(self, val):
        self.children, self.val = [], val


class PostorderIterator:
    def __init__(self, root):
        self.st = []
        self.pre = None
        node = root
        while node:
            self.st.append(node)
            if not node.children:
                break
            else:
                node = node.children[0]

    def hasNext(self):
        return len(self.st) > 0

    def next(self):
        node = self.st[-1]  # get top
        if node.children:
            next_idx = -1
            for i, child in enumerate(node.children):
                if child == self.pre:
                    next_idx = i
                    break
            if next_idx+1 != len(node.children):
                child = node.children[next_idx+1]
                while child:
                    self.st.append(child)
                    if not child.children:
                        break
                    else:
                        child = child.children[0]
        self.pre = self.st.pop()
        return self.pre

    def nextLeaf(self):
        while self.hasNext():
            node = self.next()
            if not node.children:
                return node
        return None


def isPreorderSame(r1, r2):
    it1, it2 = PostorderIterator(r1), PostorderIterator(r2)
    while it1.hasNext() and it2.hasNext():
        leaf1, leaf2 = it1.nextLeaf(), it2.nextLeaf()
        if leaf1 == leaf2 or leaf1.val == leaf2.val:
            continue
        else:
            print leaf1.val, leaf2.val
            break
    return not it1.hasNext() and not it2.hasNext()

r1 = TreeNode(1)
r1.children.extend((TreeNode(2), TreeNode(3), TreeNode(4)))
r1.children[0].children.extend((TreeNode(5), TreeNode(6)))
r2 = TreeNode(1)
r2.children.extend((TreeNode(5), TreeNode(2)))
r2.children[1].children.extend((TreeNode(6), TreeNode(3), TreeNode(4)))
print isPreorderSame(r1, r2)