# Shortest Path in a Weighted Tree **Difficulty:** HARD [External](https://leetcode.com/problems/shortest-path-in-a-weighted-tree) Canonical: https://scaleengineer.com/dsa/problems/shortest-path-in-a-weighted-tree **Algorithms:** [Depth-First Search](https://scaleengineer.com/algorithms/depth-first-search) **Data structures:** Array, Tree, Binary Indexed Tree, Segment Tree **Companies:** [Juspay](https://scaleengineer.com/companies/juspay) --- ## Problem You are given an integer `n` and an undirected, weighted tree rooted at node 1 with `n` nodes numbered from 1 to `n`. This is represented by a 2D array `edges` of length `n - 1`, where `edges[i] = [ui, vi, wi]` indicates an undirected edge from node `ui` to `vi` with weight `wi`. You are also given a 2D integer array `queries` of length `q`, where each `queries[i]` is either: * `[1, u, v, w']` – **Update** the weight of the edge between nodes `u` and `v` to `w'`, where `(u, v)` is guaranteed to be an edge present in `edges`. * `[2, x]` – **Compute** the **shortest** path distance from the root node 1 to node `x`. Return an integer array `answer`, where `answer[i]` is the **shortest** path distance from node 1 to `x` for the `ith` query of `[2, x]`. **Example 1:** **Input:** n = 2, edges = \[\[1,2,7\]\], queries = \[\[2,2\],\[1,1,2,4\],\[2,2\]\] **Output:** \[7,4\] **Explanation:** ![](https://assets.glich.co/dsa/shortest-path-in-a-weighted-tree/image0.png) * Query `[2,2]`: The shortest path from root node 1 to node 2 is 7. * Query `[1,1,2,4]`: The weight of edge `(1,2)` changes from 7 to 4. * Query `[2,2]`: The shortest path from root node 1 to node 2 is 4. **Example 2:** **Input:** n = 3, edges = \[\[1,2,2\],\[1,3,4\]\], queries = \[\[2,1\],\[2,3\],\[1,1,3,7\],\[2,2\],\[2,3\]\] **Output:** \[0,4,2,7\] **Explanation:** ![](https://assets.glich.co/dsa/shortest-path-in-a-weighted-tree/image1.png) * Query `[2,1]`: The shortest path from root node 1 to node 1 is 0. * Query `[2,3]`: The shortest path from root node 1 to node 3 is 4. * Query `[1,1,3,7]`: The weight of edge `(1,3)` changes from 4 to 7. * Query `[2,2]`: The shortest path from root node 1 to node 2 is 2. * Query `[2,3]`: The shortest path from root node 1 to node 3 is 7. **Example 3:** **Input:** n = 4, edges = \[\[1,2,2\],\[2,3,1\],\[3,4,5\]\], queries = \[\[2,4\],\[2,3\],\[1,2,3,3\],\[2,2\],\[2,3\]\] **Output:** \[8,3,2,5\] **Explanation:** ![](https://assets.glich.co/dsa/shortest-path-in-a-weighted-tree/image2.png) * Query `[2,4]`: The shortest path from root node 1 to node 4 consists of edges `(1,2)`, `(2,3)`, and `(3,4)` with weights `2 + 1 + 5 = 8`. * Query `[2,3]`: The shortest path from root node 1 to node 3 consists of edges `(1,2)` and `(2,3)` with weights `2 + 1 = 3`. * Query `[1,2,3,3]`: The weight of edge `(2,3)` changes from 1 to 3. * Query `[2,2]`: The shortest path from root node 1 to node 2 is 2. * Query `[2,3]`: The shortest path from root node 1 to node 3 consists of edges `(1,2)` and `(2,3)` with updated weights `2 + 3 = 5`. **Constraints:** * `1 <= n <= 105` * `edges.length == n - 1` * `edges[i] == [ui, vi, wi]` * `1 <= ui, vi <= n` * `1 <= wi <= 104` * The input is generated such that `edges` represents a valid tree. * `1 <= queries.length == q <= 105` * `queries[i].length == 2` or `4` * `queries[i] == [1, u, v, w']` or, * `queries[i] == [2, x]` * `1 <= u, v, x <= n` * `(u, v)` is always an edge from `edges`. * `1 <= w' <= 104` # Approaches ## Brute Force: Re-computation per Query The most straightforward approach is to simulate the process directly. For each query that asks for a shortest path distance, we can traverse the tree from the root to compute the required distance. Since the graph is a tree, the path is unique, and any standard traversal algorithm like BFS or DFS can find the distance. **Time:** O(Q * N), where Q is the number of queries and N is the number of nodes. Each distance query requires a traversal of the tree, which takes O(N) time. · **Space:** O(N) to store the adjacency list representation of the tree. **Pros:** Simple to understand and implement.; Correctly solves the problem for small inputs. **Cons:** Very high time complexity, making it infeasible for the given constraints.; Repeatedly recomputes distances for the entire tree, which is highly inefficient. ### Explanation This method involves building the graph and then handling each query as an independent event. For an update query, we modify our graph representation. A good way to represent the graph for easy updates is an adjacency list where each node maps to a hash map of its neighbors and the corresponding edge weights. This allows `O(1)` average time for finding and updating an edge weight. For a distance query, we initiate a traversal from the root, node 1. A Breadth-First Search (BFS) is a natural choice. We use a queue and a `distance` array, initialized to infinity for all nodes except the root (`distance[1] = 0`). The BFS proceeds level by level, calculating the distance to each node from the root. The distance to a neighbor `v` from a node `u` is `distance[u] + weight(u, v)`. We continue this until we have found the distance to the target node `x`, or until all nodes have been visited. ```java import java.util.*; class Solution { public long[] shortestPath(int n, int[][] edges, int[][] queries) { List> adj = new ArrayList<>(); for (int i = 0; i <= n; i++) { adj.add(new HashMap<>()); } for (int[] edge : edges) { adj.get(edge[0]).put(edge[1], edge[2]); adj.get(edge[1]).put(edge[0], edge[2]); } List answers = new ArrayList<>(); for (int[] query : queries) { if (query[0] == 1) { int u = query[1], v = query[2], w = query[3]; adj.get(u).put(v, w); adj.get(v).put(u, w); } else { int target = query[1]; answers.add(findDistance(n, adj, target)); } } long[] result = new long[answers.size()]; for (int i = 0; i < answers.size(); i++) { result[i] = answers.get(i); } return result; } private long findDistance(int n, List> adj, int target) { if (target == 1) return 0; long[] dist = new long[n + 1]; Arrays.fill(dist, -1); Queue q = new LinkedList<>(); dist[1] = 0; q.offer(1); while (!q.isEmpty()) { int u = q.poll(); if (u == target) { return dist[u]; } for (Map.Entry entry : adj.get(u).entrySet()) { int v = entry.getKey(); int weight = entry.getValue(); if (dist[v] == -1) { // In a tree, this check is sufficient if we start from root dist[v] = dist[u] + weight; q.offer(v); } } } return -1; // Should not be reached in a connected tree } } ``` ### Algorithm - Build an adjacency list representation of the tree from the `edges` array. For each edge `(u, v)` with weight `w`, add `v` to `u`'s list and `u` to `v`'s list. - For each query of type 1 `[1, u, v, w']`, update the weight of the edge `(u, v)` in the adjacency list to `w'`. - For each query of type 2 `[2, x]`, perform a graph traversal (like Breadth-First Search or Depth-First Search) starting from the root (node 1). - During the traversal, maintain an array `distance` to store the shortest distance from the root to each node. - The traversal computes the distance to all reachable nodes. The answer for the query is the computed distance to node `x`. - Collect the answers for all type 2 queries into a list. ## Precomputation with Naive Subtree Updates This approach improves upon the brute-force method by avoiding re-computation from scratch for every query. We pre-calculate all distances from the root once. When an edge weight is updated, we identify that this only affects the distances to nodes in a specific subtree. We then update the distances for only the nodes in that subtree. Querying a distance becomes a simple array lookup. **Time:** O(N + Q * N) in the worst case. Preprocessing takes O(N). Each update can take up to O(N) if the subtree is large. Each query is O(1). · **Space:** O(N) for the adjacency list, distance array, and parent array. **Pros:** Reduces query time to O(1).; More efficient than pure brute-force if updates affect small subtrees. **Cons:** The update operation can take O(N) time in the worst case (e.g., updating an edge connected to the root), leading to the same worst-case time complexity as the brute-force approach.; Inefficient for test cases with many updates on edges high up in the tree. ### Explanation First, we perform a one-time traversal (DFS or BFS) from the root (node 1) to populate a `distance` array with the initial shortest path distances to all other nodes. We also determine the parent of each node, which is crucial for identifying subtrees. When an update query for an edge `(u, v)` arrives, we first determine the parent-child relationship. Let's say `u` is the parent of `v`. A change in the weight of edge `(u, v)` affects the path distance from the root to `v` and to all of `v`'s descendants. The change in distance is uniform for all these nodes. So, we calculate the difference (`delta`) between the new and old weights and then start another traversal from `v` downwards, updating the `distance` for each node in `v`'s subtree by adding `delta`. Distance queries are now very fast, as we just need to look up the pre-calculated (and updated) distance in our `distance` array. ```java // Assume distance and parent arrays are pre-filled. // adj is a List> // Inside the main query processing loop: if (query[0] == 1) { int u = query[1], v = query[2], w_new = query[3]; // Ensure u is the parent of v if (parent[u] == v) { int temp = u; u = v; v = temp; } long w_old = adj.get(u).get(v); long delta = w_new - w_old; adj.get(u).put(v, w_new); adj.get(v).put(u, w_new); // Update subtree of v Queue q = new LinkedList<>(); q.offer(v); distance[v] += delta; while (!q.isEmpty()) { int curr = q.poll(); for (int neighbor : adj.get(curr).keySet()) { if (neighbor != parent[curr]) { // Go down the tree distance[neighbor] += delta; q.offer(neighbor); } } } } else { // query[0] == 2 int x = query[1]; answers.add(distance[x]); } ``` This snippet shows the logic for a single query. The full implementation would require the initial DFS/BFS to set up `distance` and `parent` arrays. ### Algorithm - **Preprocessing:** - Build an adjacency list for the tree. - Perform a single DFS/BFS traversal from the root (node 1) to compute initial distances to all nodes and store them in a `distance` array. Also, compute the parent of each node in the tree rooted at 1. - **Query Processing:** - For an update query `[1, u, v, w']`: - Determine which node is the parent (e.g., `u` is parent of `v`). - Find the old weight of the edge `(u, v)` and calculate the change in weight, `delta = w' - w_old`. - Update the weight in the adjacency list. - The distances for all nodes in the subtree of `v` are affected by this change. Traverse the subtree of `v` (e.g., using DFS/BFS starting from `v`) and add `delta` to the stored distance of each node in the subtree. - For a distance query `[2, x]`, simply return the value from the `distance[x]` array. ## Efficient Approach: Fenwick Tree on Flattened Tree The key observation for an efficient solution is that an edge weight update affects all nodes in a specific subtree uniformly. This pattern of "range updates" (on a subtree) and "point queries" (for a node's distance) suggests using a specialized data structure. By linearizing the tree using DFS start and end times, we can map the subtree updates to range updates on an array. A Fenwick Tree (or a Segment Tree) is an excellent tool for handling these operations efficiently. **Time:** O(N + Q * log N). Preprocessing (DFS) takes O(N). Each of the Q queries takes O(log N) for BIT operations. · **Space:** O(N) for the adjacency list, parent/distance/time arrays, and the Fenwick Tree. **Pros:** Highly efficient, with logarithmic time complexity for both updates and queries.; Scales well for large inputs, passing the given constraints.; It's a standard and powerful technique for a class of problems involving queries on trees. **Cons:** More complex to implement due to the need for DFS traversal times and a Fenwick Tree.; The constant factors might be higher than simpler approaches for very small N. ### Explanation This approach combines tree algorithms with a data structure to handle the queries efficiently. **1. Preprocessing:** We start with a DFS from the root (node 1). During this traversal, we compute several properties for each node: its parent, its initial distance from the root, and its DFS start and end times. The start time is recorded when a node is first visited, and the end time is recorded after all its descendants have been visited. This mapping ensures that for any node `v`, all nodes `u` in its subtree satisfy `startTime[v] <= startTime[u] <= endTime[v]`. **2. Fenwick Tree for Updates and Queries:** We use a Fenwick Tree (BIT) to manage the distance modifications. A standard BIT supports point updates and prefix sum queries. To handle range updates and point queries, we can use a clever trick: to add a value `delta` to a range `[l, r]`, we perform two point updates on the BIT: `add(l, delta)` and `add(r + 1, -delta)`. The effect of this is that when we query the prefix sum up to an index `i` (`query(i)`), we get the sum of all deltas for ranges that start at or before `i`. This is exactly the total change affecting the node corresponding to time `i`. **3. Handling Queries:** - **Update `[1, u, v, w']`:** We find the child node (say `v`), calculate `delta`, and update the BIT at `startTime[v]` and `endTime[v] + 1`. This takes `O(log N)` time. - **Query `[2, x]`:** The current distance is the initial distance plus all accumulated changes. We retrieve this by querying the BIT at `startTime[x]`. The result is `initial_dist[x] + bit.query(startTime[x])`. This also takes `O(log N)` time. ```java import java.util.*; class Solution { private int timer; private int[] parent, startTime, endTime; private long[] initialDist; private List> adj; public long[] shortestPath(int n, int[][] edges, int[][] queries) { adj = new ArrayList<>(); for (int i = 0; i <= n; i++) adj.add(new HashMap<>()); for (int[] edge : edges) { adj.get(edge[0]).put(edge[1], edge[2]); adj.get(edge[1]).put(edge[0], edge[2]); } parent = new int[n + 1]; startTime = new int[n + 1]; endTime = new int[n + 1]; initialDist = new long[n + 1]; timer = 0; dfs(1, 0, 0); FenwickTree bit = new FenwickTree(n); List answers = new ArrayList<>(); for (int[] query : queries) { if (query[0] == 1) { int u = query[1], v = query[2], w = query[3]; if (parent[u] == v) { // Ensure u is parent of v int temp = u; u = v; v = temp; } long oldWeight = adj.get(u).get(v); long delta = w - oldWeight; adj.get(u).put(v, w); adj.get(v).put(u, w); bit.add(startTime[v], delta); bit.add(endTime[v] + 1, -delta); } else { int x = query[1]; long currentChange = bit.query(startTime[x]); answers.add(initialDist[x] + currentChange); } } long[] result = new long[answers.size()]; for (int i = 0; i < answers.size(); i++) result[i] = answers.get(i); return result; } private void dfs(int u, int p, long currentDist) { parent[u] = p; initialDist[u] = currentDist; startTime[u] = ++timer; for (Map.Entry entry : adj.get(u).entrySet()) { int v = entry.getKey(); if (v != p) { dfs(v, u, currentDist + entry.getValue()); } } endTime[u] = timer; } } class FenwickTree { private long[] bit; private int size; public FenwickTree(int n) { this.size = n; this.bit = new long[n + 2]; } public void add(int index, long delta) { for (; index <= size; index += index & -index) { bit[index] += delta; } } public long query(int index) { long sum = 0; for (; index > 0; index -= index & -index) { sum += bit[index]; } return sum; } } ``` ### Algorithm - **Preprocessing:** - Build an adjacency list. - Perform a DFS from the root (node 1) to compute: - `initial_dist[i]`: The initial distance from the root to node `i`. - `parent[i]`: The parent of node `i`. - `startTime[i]` and `endTime[i]`: The entry and exit times for each node in the DFS traversal. This linearizes the tree such that all nodes in a subtree of `v` have start times in the range `[startTime[v], endTime[v]]`. - **Data Structure:** - Initialize a Fenwick Tree (BIT) of size `N` with all zeros. This BIT will store the accumulated delta changes. - **Query Processing:** - For an update query `[1, u, v, w']`: - Determine the child node (e.g., `v`). - Calculate `delta = w' - w_old`. - This `delta` applies to all nodes in the subtree of `v`. In our linearized representation, this corresponds to the range of start times `[startTime[v], endTime[v]]`. - Perform a range update on the BIT: `bit.add(startTime[v], delta)` and `bit.add(endTime[v] + 1, -delta)`. - For a distance query `[2, x]`: - The total change in distance for node `x` is the sum of all deltas for ranges that include `startTime[x]`. This can be found with a point query on our BIT structure, which is `bit.query(startTime[x])`. - The final distance is `initial_dist[x] + bit.query(startTime[x])`.