剑指offer python版 34.二叉树中和为某一值的路径

"""
输入一棵二叉树和整数，打印出二叉树中节点值的和为
输入整数的所有路径
"""
class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

class Solution:
    def FindPath(self, root, sum):
        if root == None:
            return []
        if root.left == None and root.right == None:
            if sum == root.val:
                return [[root.val]]
            else:
                return []

        stack = []
        leftStack = self.pathSum(root.left, sum-root.val)
        for i in leftStack:
            i.insert(0, root.val)
            stack.append(i)
        rightStack = self.pathSum(root.right, sum-root.val)
        for i in rightStack:
            i.insert(0, root.val)
            stack.append(i)


    def pathSum(self, root, sum):
        if root == None:
            return []
        if root.left == None and root.right == None:
            if sum == root.val:
                return [[root.val]]
            else:
                return []
        a = self.pathSum(root.left, sum - root.val) + self.pathSum(root.right, sum - root.val)
        return [[root.val] + i for i in a]

    def FindPath2(self, root, sum):
        if root == None:
            return []
        if root.left == None and root.right == None and root.val == sum:
            return [[root.val]]

        result = []
        left = self.FindPath2(root.left, sum-root.val)
        right = self.FindPath2(root.right, sum-root.val)
        for i in left+right:
            result.append([root.val] + i)
        return result

if __name__ == '__main__':
    pNode1 = TreeNode(10)
    pNode2 = TreeNode(5)
    pNode3 = TreeNode(12)
    pNode4 = TreeNode(4)
    pNode5 = TreeNode(7)

    pNode1.left = pNode2
    pNode1.right = pNode3
    pNode2.left = pNode4
    pNode2.right = pNode5

    S = Solution()
    print(S.FindPath(pNode1, 22))
    print(S.FindPath2(pNode1, 22))