Depth-First Search (DFS) II

Another problem whose a solution can be accomplished via Depth-First Search (DFS). Map the tree to a hash table. Calculate the height of the tree separately (also a DFS). Then perform a DFS to calculate the weighted sum. Code is down below, cheers, ACC.

Weighted Sum of a Tree - LeetCode

You are given an integer array parent of length n representing a rooted tree with nodes labeled from 0 to n - 1.

The tree is rooted at node 0, so parent[0] = -1. For each node i where 1 <= i <= n - 1, parent[i] denotes the parent of node i.

You are also given an integer array nums of length n, where nums[i] denotes the value of node i.

The weight of a node i at depth d is nums[i] * (h - d + 1), where h is the height of the tree.

Return the sum of the weights of all nodes in the tree.

The depth of a node is the number of nodes on the path from the root to that node, inclusive, with the root having depth 1.

The height of the tree is the maximum depth among all nodes in the tree.

 

Example 1:

​​​​​​​

Input: parent = [-1,0,0,0,2,2], nums = [5,2,3,1,4,6]

Output: 37

Explanation:

The height of the tree is 3.

Node	nums[i]	Depth (d)	Weight
0	5	1	5 * (3 - 1 + 1) = 15
1	2	2	2 * (3 - 2 + 1) = 4
2	3	2	3 * (3 - 2 + 1) = 6
3	1	2	1 * (3 - 2 + 1) = 2
4	4	3	4 * (3 - 3 + 1) = 4
5	6	3	6 * (3 - 3 + 1) = 6

The sum of all node weights is 15 + 4 + 6 + 2 + 4 + 6 = 37.

Example 2:

​​​​​​​​​​​​​​

Input: parent = [-1,0,1,2], nums = [1,2,3,4]

Output: 20

Explanation:

The height of the tree is 4.

Node	nums[i]	Depth (d)	Weight
0	1	1	1 * (4 - 1 + 1) = 4
1	2	2	2 * (4 - 2 + 1) = 6
2	3	3	3 * (4 - 3 + 1) = 6
3	4	4	4 * (4 - 4 + 1) = 4

The sum of all node weights is 4 + 6 + 6 + 4 = 20.

 

Constraints:

• 1 <= n <= 105
• n == parent.length == nums.length
• parent[0] == -1
• 0 <= parent[i] <= n - 1 for all i in [1, n - 1]
• 1 <= nums[i] <= 106
• The input is generated such that the array parent represents a valid tree rooted at node 0.

public class Solution {
    public long WeightedSum(int[] parent, int[] nums)
    {
        Hashtable tree = new Hashtable();

        for (int child = 0; child < parent.Length; child++)
        {
            if (!tree.ContainsKey(parent[child])) tree.Add(parent[child], new Hashtable());
            Hashtable children = (Hashtable)tree[parent[child]];
            children.Add(child, nums[child]);
        }

        int height = CalculateHeightTree(tree, -1);

        long retVal = 0;
        _WeightedSum(tree, -1, 0, height, ref retVal);

        return retVal;
    }

    private void _WeightedSum(Hashtable tree, int node, int level, int height, ref long retVal)
    {
        if (!tree.Contains(node)) return;

        Hashtable children = (Hashtable)tree[node];
        foreach (int child in children.Keys)
        {
            retVal += ((int)children[child] * (height - (level + 1) + 1L));
            _WeightedSum(tree, child, level + 1, height, ref retVal);
        }
    }

    private int CalculateHeightTree(Hashtable tree, int node)
    {
        if (!tree.ContainsKey(node)) return 0;

        int height = 0;
        Hashtable children = (Hashtable)tree[node];
        foreach (int n in children.Keys)
        {
            height = Math.Max(height, CalculateHeightTree(tree, n));
        }
        return height + 1;
    }
}