Katkat
Back

Build a Tree Whose Nodes Hold the Original Tree's Subtree Sums

SnowflakeSoftware Engineer · Technical Screen
MediumSoftware EngineerTechnical ScreenCoding & Algorithms

Last updated

Given a rooted tree as a parent array with a value on every node, return a same-shaped tree whose nodes hold the sum of each original subtree. It tests aggregating children before parents without hitting recursion limits, and leaving the original tree unchanged.

Every node of a rooted tree holds an integer. Build a second tree with the same shape in which each node holds the sum of all the values in the corresponding subtree of the original tree, leaving the original tree as it is. The new tree has the same parent links as the original, so return only its values.

Function Signature

def subtree_sum_tree(parent: list[int], values: list[int]) -> list[int]:

Rules

  • The tree has n nodes numbered 0 to n - 1. Node 0 is the root and parent[0] == -1. For every other node i, parent[i] is the index of its parent. A node can have any number of children, and parents can have larger indices than their children.
  • values[i] is the value held by node i in the original tree.
  • The subtree of a node consists of the node and all of its descendants.
  • Return a list sums of length n in which sums[i] is the total of values[j] over every node j in the subtree of node i. Together with parent, sums describes the new tree.

Constraints

  • 1 <= n == len(parent) == len(values) <= 10^5
  • parent[0] == -1, and for i >= 1, 0 <= parent[i] <= n - 1; the parent links form a single tree rooted at node 0 with no cycles.
  • -10^4 <= values[i] <= 10^4, so every sum lies between -10^9 and 10^9.
  • The tree can be a single path, so its depth can reach n.

Examples

Example 1

 0
 / \
 1 2
 / \ \
 3 4 6
 |
 5

Input: parent = [-1, 0, 0, 1, 1, 3, 2], values = [5, 3, -2, 4, 1, 6, 7]
Output: [24, 14, 5, 10, 1, 6, 7]

The drawing shows node indices. Node 3's subtree holds nodes 3 and 5, so its sum is 4 + 6 = 10. Node 1's subtree holds nodes 1, 3, 4 and 5: 3 + 4 + 1 + 6 = 14. Node 2's sum is -2 + 7 = 5, and the root's is 5 + 14 + 5 = 24. Leaves keep their own values.

Example 2

Input: parent = [-1, 2, 0, 1], values = [1, -1, 1, -1]
Output: [0, -2, -1, -1]

The tree is the path 0, 2, 1, 3. Node 1's subtree holds nodes 1 and 3 (sum -2), node 2's holds nodes 2, 1 and 3 (sum -1), and the root's holds all four nodes (sum 0).

Example 3

Input: parent = [-1], values = [8]
Output: [8]