Skip to content

4015. Weighted Sum of a Tree

Description

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.

Solutions

Solution 1: BFS

Thinking

The weight of node \(i\) is \(\textit{nums}[i]\times(h-d_i+1)\). Computing the height first and then summing by definition needs two traversals and stored depths.

Splitting the sum into \(h\sum\textit{nums}[i]+\sum\textit{nums}[i](1-d_i)\) lets one BFS accumulate the second term while it walks; the number of layers at the end is \(h\).

The adjacency lists keep only parent-to-child edges, so the level order matches the depth definition.

The weight of node \(i\) is \(\textit{nums}[i] \times (h - d_i + 1)\), where \(d_i\) is the depth of node \(i\) and \(h\) is the height of the tree. Therefore, the sum of the weights of all nodes is:

\[\sum_{i=0}^{n-1} \textit{nums}[i] \times (h - d_i + 1) = h \times \sum_{i=0}^{n-1} \textit{nums}[i] + \sum_{i=0}^{n-1} \textit{nums}[i] \times (1 - d_i)\]

We can use BFS to traverse the tree level by level. During the traversal, we maintain the current level \(d\) (the root is at level \(1\)) and accumulate \(\textit{nums}[i] \times (1 - d)\) for each node. After the traversal finishes, \(d\) equals the height \(h\) of the tree, and adding \(h \times \sum \textit{nums}[i]\) gives the answer.

The time complexity is \(O(n)\), and the space complexity is \(O(n)\), where \(n\) is the number of nodes.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
class Solution:
    def weightedSum(self, parent: list[int], nums: list[int]) -> int:
        n = len(nums)
        g = [[] for _ in range(n)]
        for i in range(1, n):
            g[parent[i]].append(i)
        ans = 0
        q = [0]
        d = 0
        while q:
            d += 1
            nq = []
            for i in q:
                ans += nums[i] * (1 - d)
                nq.extend(g[i])
            q = nq
        ans += d * sum(nums)
        return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class Solution {
    public long weightedSum(int[] parent, int[] nums) {
        int n = nums.length;

        List<Integer>[] g = new ArrayList[n];
        Arrays.setAll(g, e -> new ArrayList<>());

        for (int i = 1; i < n; i++) {
            g[parent[i]].add(i);
        }

        long ans = 0;

        List<Integer> q = new ArrayList<>();
        q.add(0);

        int d = 0;

        while (!q.isEmpty()) {
            d++;

            List<Integer> nq = new ArrayList<>();

            for (int i : q) {
                ans += (long) nums[i] * (1 - d);
                nq.addAll(g[i]);
            }

            q = nq;
        }

        long sum = 0;
        for (int x : nums) {
            sum += x;
        }

        ans += (long) d * sum;

        return ans;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
class Solution {
public:
    long long weightedSum(vector<int>& parent, vector<int>& nums) {
        int n = nums.size();

        vector<vector<int>> g(n);

        for (int i = 1; i < n; i++) {
            g[parent[i]].push_back(i);
        }

        long long ans = 0;

        vector<int> q = {0};

        int d = 0;

        while (!q.empty()) {
            d++;

            vector<int> nq;

            for (int i : q) {
                ans += 1LL * nums[i] * (1 - d);
                for (int son : g[i]) {
                    nq.push_back(son);
                }
            }

            q = move(nq);
        }

        long long sum = 0;
        for (int x : nums) {
            sum += x;
        }

        ans += 1LL * d * sum;

        return ans;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
func weightedSum(parent []int, nums []int) int64 {
    n := len(nums)

    g := make([][]int, n)

    for i := 1; i < n; i++ {
        g[parent[i]] = append(g[parent[i]], i)
    }

    var ans int64

    q := []int{0}

    d := 0

    for len(q) > 0 {
        d++

        nq := make([]int, 0)

        for _, i := range q {
            ans += int64(nums[i]) * int64(1-d)

            for _, son := range g[i] {
                nq = append(nq, son)
            }
        }

        q = nq
    }

    var sum int64
    for _, x := range nums {
        sum += int64(x)
    }

    ans += int64(d) * sum

    return ans
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
function weightedSum(parent: number[], nums: number[]): number {
    const n = nums.length;

    const g: number[][] = Array.from({ length: n }, () => []);

    for (let i = 1; i < n; i++) {
        g[parent[i]].push(i);
    }

    let ans = 0;

    let q: number[] = [0];

    let d = 0;

    while (q.length > 0) {
        d++;

        const nq: number[] = [];

        for (const i of q) {
            ans += nums[i] * (1 - d);

            for (const son of g[i]) {
                nq.push(son);
            }
        }

        q = nq;
    }

    let sum = 0;
    for (const x of nums) {
        sum += x;
    }

    ans += d * sum;

    return ans;
}

Comments