666. Path Sum IV

Information

  • Diffculty: Medium

  • Created: 2020-02-11 17:55:00

  • Last Motified: 2020-02-12 20:02:00

  • Tags: Tree

Intuition

The depth of the tree is smaller than 5, so there are at most 2 ^ 4 - 1 = 15 nodes in the tree. So simply we could recover this tree from the “list num” representation to a general tree representation. Then, we could easily get sum of all path from leaves to root, by doing preorder traversal and update each node’s value to the sum of itself and its parent. Finally, for each leave, add the value to the total_sum.

Solution

class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None


class Solution:
    def pathSum(self, nums: List[int]) -> int:
        if not nums:
            return 0
        nums = [[int(s[0]) - 1, int(s[1]) - 1, int(s[2])] for s in map(str, nums)]
        n = len(nums)
        nodes = [TreeNode(v[-1]) for v in nums]

        lo, hi = 0, 1
        while hi < n:
            if nums[lo][1] == nums[hi][1] // 2 and nums[lo][0] == nums[hi][0] - 1:
                if 2 * nums[lo][1] == nums[hi][1]:
                    nodes[lo].left = nodes[hi]
                else:
                    nodes[lo].right = nodes[hi]
                hi += 1
            else:
                lo += 1

        return self._dfs(nodes[0])

    def _dfs(self, node, previous=0):
        if node is None:
            return 0

        current = previous + node.val
        if node.left is None and node.right is None:
            # The current node is a leaf and we could end the DFS
            return current
        else:
            return self._dfs(node.left, current) + self._dfs(node.right, current)