Quick Overview

Given an undirected tree whose edges carry positive integer weights, return for every node the sum of its weighted distances to all other nodes. The problem tests tree traversal, subtree aggregation and producing all n answers fast enough for trees with tens of thousands of nodes, with sums that need 64-bit integers.

Weighted Tree: Total Distance from Each Node to All Other Nodes

Company: Google

Role: Software Engineer

Category: Coding & Algorithms

Difficulty: hard

Interview Round: Onsite

An undirected tree has `n` nodes labeled `0` to `n - 1`, and every edge has a positive integer weight. The distance between two nodes is the sum of the edge weights on the unique path that joins them. Write a function that returns, for every node, the sum of its distances to all the other nodes. ### Function Signature ```python def weighted_distance_sums(n: int, edges: list[list[int]]) -> list[int]: ``` ### Rules - `edges[i] = [u, v, w]` is an undirected edge of weight `w` between nodes `u` and `v`. `[u, v, w]` and `[v, u, w]` describe the same edge, and each edge is listed exactly once. - The edges always form one tree: there are exactly `n - 1` of them, and every pair of nodes is joined by exactly one path. - Return a list `answer` of length `n` in which `answer[i]` is the sum, over every node `j` other than `i`, of the distance between `i` and `j`. The list is indexed by node label, so the answer is unique. - When `n = 1` there are no edges, and the answer is `[0]`. ### Constraints - `1 <= n <= 3 * 10^4` - `len(edges) == n - 1` - `0 <= u, v <= n - 1` and `u != v` for every edge - `1 <= w <= 10^4` for every edge - `edges` is in no particular order. - A single `answer[i]` can reach about `4.5 * 10^12`, which exceeds `2^31 - 1`: use 64-bit integers in a language with fixed-width integers. Every value stays below `2^53`. ### Examples **Example 1** ```text Input: n = 6, edges = [[0, 1, 2], [0, 2, 1], [2, 3, 3], [2, 4, 1], [2, 5, 2]] Output: [12, 20, 10, 22, 14, 18] ``` From node `0`, the distances to nodes `1`, `2`, `3`, `4` and `5` are `2`, `1`, `4`, `2` and `3`, which sum to `12`. From node `3`, the distances to nodes `0`, `1`, `2`, `4` and `5` are `4`, `6`, `3`, `4` and `5`, which sum to `22`. **Example 2** ```text Input: n = 3, edges = [[1, 0, 5], [2, 1, 1]] Output: [11, 6, 7] ``` The tree is the path `0 - 1 - 2`. Node `0` is at distance `5` from node `1` and `6` from node `2`; node `1` is at distance `5` from node `0` and `1` from node `2`; node `2` is at distance `6` from node `0` and `1` from node `1`. **Example 3** ```text Input: n = 1, edges = [] Output: [0] ``` A single node has no other nodes to reach.

Overview: Given an undirected tree whose edges carry positive integer weights, return for every node the sum of its weighted distances to all other nodes. The problem tests tree traversal, subtree aggregation and producing all n answers fast enough for trees with tens of thousands of nodes, with sums that need 64-bit integers.

An undirected tree has `n` nodes labeled `0` to `n - 1`, and every edge has a positive integer weight. The distance between two nodes is the sum of the edge weights on the unique path that joins them. Return, for every node, the sum of its distances to all the other nodes. Implement `weighted_distance_sums(n, edges)`. ### Rules - `edges[i] = [u, v, w]` is an undirected edge of weight `w` between nodes `u` and `v`. `[u, v, w]` and `[v, u, w]` describe the same edge, and each edge is listed exactly once. - The edges always form one tree: there are exactly `n - 1` of them, and every pair of nodes is joined by exactly one path. - Return a list `answer` of length `n` in which `answer[i]` is the sum, over every node `j` other than `i`, of the distance between `i` and `j`. The list is indexed by node label (position `i` holds the sum for node `i`), so the answer is unique. - When `n = 1` there are no edges, and the answer is `[0]`. ### Example 1 ```text Input: n = 6, edges = [[0, 1, 2], [0, 2, 1], [2, 3, 3], [2, 4, 1], [2, 5, 2]] Output: [12, 20, 10, 22, 14, 18] ``` From node `0`, the distances to nodes `1`, `2`, `3`, `4` and `5` are `2`, `1`, `4`, `2` and `3`, which sum to `12`. From node `3`, the distances to nodes `0`, `1`, `2`, `4` and `5` are `4`, `6`, `3`, `4` and `5`, which sum to `22`. ### Example 2 ```text Input: n = 3, edges = [[1, 0, 5], [2, 1, 1]] Output: [11, 6, 7] ``` The tree is the path `0 - 1 - 2`. Node `0` is at distance `5` from node `1` and `6` from node `2`; node `1` is at distance `5` from node `0` and `1` from node `2`; node `2` is at distance `6` from node `0` and `1` from node `1`. ### Constraints - `1 <= n <= 3 * 10^4` - `len(edges) == n - 1` - `0 <= u, v <= n - 1` and `u != v` for every edge - `1 <= w <= 10^4` for every edge - `edges` is in no particular order. - A single `answer[i]` can reach about `4.5 * 10^12`, which exceeds `2^31 - 1`: use 64-bit integers in a language with fixed-width integers (`long` in Java, `long long` in C++). Every value stays below `2^53`.

Constraints

  • 1 <= n <= 3 * 10^4
  • len(edges) == n - 1
  • 0 <= u, v <= n - 1 and u != v for every edge
  • 1 <= w <= 10^4 for every edge
  • edges is in no particular order; [u, v, w] and [v, u, w] describe the same edge, and each edge is listed exactly once
  • The edges always form one tree; when n = 1 there are no edges and the answer is [0]
  • A single answer[i] can reach about 4.5 * 10^12, which exceeds 2^31 - 1: use 64-bit integers (long in Java, long long in C++). Every value stays below 2^53.

Examples

Input: (6, [[0, 1, 2], [0, 2, 1], [2, 3, 3], [2, 4, 1], [2, 5, 2]])

Expected Output: [12, 20, 10, 22, 14, 18]

Explanation: Source Example 1: node 0 sums 2+1+4+2+3 = 12, node 3 sums 4+6+3+4+5 = 22; mixed weights, so a count-only (unweighted) sum gives a different list.

Input: (3, [[1, 0, 5], [2, 1, 1]])

Expected Output: [11, 6, 7]

Explanation: Source Example 2: path 0 - 1 - 2 with every edge given as [child, parent]; a solution that treats edges as directed from u cannot reach node 1 from node 0.

Hints

  1. The input names no root and fixes no edge order or orientation, so store each edge in both directions and choose a starting node yourself.
  2. With up to 3 * 10^4 nodes, a separate traversal from every node repeats almost all of the work. Two neighbouring nodes see nearly the same distances: crossing one edge of weight w changes every other node's distance by exactly w, up or down.
  3. Sums can exceed 2^31 - 1, so keep running totals in 64-bit integers in Java and C++.

Community answers

Answer by ahiremath22

Brute force: DFS for every node - o(n^2) Apply rerooting algo once you calculate the weight from say root 0 and subtree size from root 0. Once this is calculated, move the node and for each re root, subtract the subtree size -

Loading coding console...

Show the approach

Approach

Algorithm (rerooting in two linear passes). Store every edge in both directions, since an edge may be listed as [u, v, w] or [v, u, w] and in any order. Root the tree at node 0 and run an iterative depth-first search (no recursion, so a 3 * 10^4-node path is safe); it records each node's parent, the weight of the edge to that parent, and a visit order in which every parent comes before its children.

Pass 1 walks that order backwards (children before parents). For each node x it keeps size[x], the number of nodes in x's subtree, and down[x], the sum of distances from x to the nodes in its subtree. A child c joined to p by weight w adds size[c] to size[p] and down[c] + w * size[c] to down[p], because each of the size[c] nodes under c is w farther from p than from c. After the pass, answer[0] = down[0].

Pass 2 walks the order forwards (parents before children). Moving the reference point from p to its child c across an edge of weight w brings the size[c] nodes of c's subtree exactly w closer and pushes the other n - size[c] nodes exactly w farther, so answer[c] = answer[p] + w * (n - 2 * size[c]).

Invariant and correctness: when pass 2 reaches c, answer[p] is already final because p precedes c in the order. The path between two nodes crosses edge (p, c) exactly when one endpoint lies in c's subtree and the other does not, so the shift above accounts for every pair, and induction from the root gives every answer[i] exactly.

Edge cases: n = 1 has no edges, the search visits only node 0, and the answer is [0]. Node 0 may be a leaf or a hub; the edge order and orientation never matter because the parent links come from the search, not from the input. Sums reach about 4.5 * 10^12, so Java uses long and C++ long long; JavaScript numbers stay exact because every value is below 2^53.

Time complexity:
O(n)
Space complexity:
O(n)