leetcode - 666. Path Sum IV

Description

If the depth of a tree is smaller than 5, then this tree can be represented by an array of three-digit integers. For each integer in this array:

• The hundreds digit represents the depth d of this node where 1 <= d <= 4.
• The tens digit represents the position p of this node in the level it belongs to where 1 <= p <= 8. The position is the same as that in a full binary tree.
• The units digit represents the value v of this node where 0 <= v <= 9.

Given an array of ascending three-digit integers nums representing a binary tree with a depth smaller than 5, return the sum of all paths from the root towards the leaves.

It is guaranteed that the given array represents a valid connected binary tree.

Example 1:

Input: nums = [113,215,221]
Output: 12
Explanation: The tree that the list represents is shown.
The path sum is (3 + 5) + (3 + 1) = 12.

Example 2:

Input: nums = [113,221]
Output: 4
Explanation: The tree that the list represents is shown. 
The path sum is (3 + 1) = 4.

Constraints:

1 <= nums.length <= 15
110 <= nums[i] <= 489
nums represents a valid binary tree with depth less than 5

Solution

For each node, the path sum is: left_sum + right_sum + val * number_of_leaves. So we use a hashmap to store such information, and go through the right to the left (like post-order traverse)

Time complexity: $o(n)$
Space complexity: $o(n)$

Code

class Solution:
    def pathSum(self, nums: List[int]) -> int:
        if len(nums) == 1:
            return nums[0] % 10
        path_sum = {}
        for i in range(len(nums) - 1, -1, -1):
            level = nums[i] // 100
            position = (nums[i] - level * 100) // 10
            val = nums[i] % 10
            parent_key = (level - 1) * 10 + (position + 1) // 2
            cur_key = level * 10 + position
            if parent_key not in path_sum:
                path_sum[parent_key] = [0, 0, 0]
            if cur_key not in path_sum:
                path_sum[cur_key] = [0, 0, 0]
            if sum(path_sum[cur_key]) == 0:
                # leaf node
                cur_sum = val
                cur_child_cnt = 1
            else:
                # not a leaf node
                cur_child_cnt = path_sum[cur_key][2]
                cur_sum = path_sum[cur_key][0] + path_sum[cur_key][1] + val * cur_child_cnt
            # even
            if position & 1 == 0:
                path_sum[parent_key][1] += cur_sum
                path_sum[parent_key][2] += cur_child_cnt
            else:
                path_sum[parent_key][0] += cur_sum
                path_sum[parent_key][2] += cur_child_cnt
        return path_sum[11][0] + path_sum[11][1] + val * path_sum[11][2]