Minimum Edge Weight Equilibrium Queries in a Tree

Hard

#2545Time: O(M * N), where M is the number of queries and N is the number of nodes. For each of the M queries, we perform a traversal (BFS/DFS) which takes O(N) time in a tree.Space: O(N), where N is the number of nodes. This is for storing the adjacency list, the parent array, and the queue for BFS for each query.1 company

Data structures

Array Tree Graph

Companies

Sprinklr

Prompt

There is an undirected tree with n nodes labeled from 0 to n - 1. You are given the integer n and a 2D integer array edges of length n - 1, where edges[i] = [u_i, v_i, w_i] indicates that there is an edge between nodes u_i and v_i with weight w_i in the tree.

You are also given a 2D integer array queries of length m, where queries[i] = [a_i, b_i]. For each query, find the minimum number of operations required to make the weight of every edge on the path from a_i to b_i equal. In one operation, you can choose any edge of the tree and change its weight to any value.

Note that:

Queries are independent of each other, meaning that the tree returns to its initial state on each new query.
The path from a_i to b_i is a sequence of distinct nodes starting with node a_i and ending with node b_i such that every two adjacent nodes in the sequence share an edge in the tree.

Return an array answer of length m where answer[i] is the answer to the i^th query.

Example 1:

Input: n = 7, edges = [[0,1,1],[1,2,1],[2,3,1],[3,4,2],[4,5,2],[5,6,2]], queries = [[0,3],[3,6],[2,6],[0,6]]
Output: [0,0,1,3]
Explanation: In the first query, all the edges in the path from 0 to 3 have a weight of 1. Hence, the answer is 0.
In the second query, all the edges in the path from 3 to 6 have a weight of 2. Hence, the answer is 0.
In the third query, we change the weight of edge [2,3] to 2. After this operation, all the edges in the path from 2 to 6 have a weight of 2. Hence, the answer is 1.
In the fourth query, we change the weights of edges [0,1], [1,2] and [2,3] to 2. After these operations, all the edges in the path from 0 to 6 have a weight of 2. Hence, the answer is 3.
For each queries[i], it can be shown that answer[i] is the minimum number of operations needed to equalize all the edge weights in the path from ai to bi.

Example 2:

Input: n = 8, edges = [[1,2,6],[1,3,4],[2,4,6],[2,5,3],[3,6,6],[3,0,8],[7,0,2]], queries = [[4,6],[0,4],[6,5],[7,4]]
Output: [1,2,2,3]
Explanation: In the first query, we change the weight of edge [1,3] to 6. After this operation, all the edges in the path from 4 to 6 have a weight of 6. Hence, the answer is 1.
In the second query, we change the weight of edges [0,3] and [3,1] to 6. After these operations, all the edges in the path from 0 to 4 have a weight of 6. Hence, the answer is 2.
In the third query, we change the weight of edges [1,3] and [5,2] to 6. After these operations, all the edges in the path from 6 to 5 have a weight of 6. Hence, the answer is 2.
In the fourth query, we change the weights of edges [0,7], [0,3] and [1,3] to 6. After these operations, all the edges in the path from 7 to 4 have a weight of 6. Hence, the answer is 3.
For each queries[i], it can be shown that answer[i] is the minimum number of operations needed to equalize all the edge weights in the path from ai to bi.

Constraints:

1 <= n <= 10⁴
edges.length == n - 1
edges[i].length == 3
0 <= u_i, v_i < n
1 <= w_i <= 26
The input is generated such that edges represents a valid tree.
1 <= queries.length == m <= 2 * 10⁴
queries[i].length == 2
0 <= a_i, b_i < n

Approaches

2 approaches with complexity analysis and trade-offs.

This approach handles each query independently by performing a graph traversal to find the path between the two given nodes. For each query, it finds the path, collects the weights of all edges on that path, and then calculates the answer.

Algorithm

For each query (a, b):
1. Build an adjacency list representation of the tree if not already built. The list should store pairs of (neighbor, weight).
2. Perform a Breadth-First Search (BFS) starting from node a to find the unique path to node b. During the BFS, maintain a parent array to reconstruct the path later. parent[i] will store the node from which we reached node i.
3. Once the BFS reaches b, the search can stop. The path is now implicitly stored in the parent array.
4. Reconstruct the path by backtracking from b to a using the parent array. While backtracking from a node curr to its parent p = parent[curr], find the weight of the edge (curr, p).
5. Store the frequencies of all edge weights encountered on the path in a hash map or an array.
6. Keep track of the total number of edges on the path (pathLength).
7. After collecting all weights, find the maximum frequency (maxFreq) among them.
8. The minimum number of operations for the query is pathLength - maxFreq.

Walkthrough

The core idea is to solve each query from scratch. Given a query with nodes a and b, we need to find the path connecting them. In a tree, this path is unique. We can use a graph traversal algorithm like Breadth-First Search (BFS) or Depth-First Search (DFS) to find this path. A straightforward way is to start a BFS from a, keeping track of the parent of each visited node. When we reach b, we can trace back the path to a using the parent pointers.

As we trace the path, we collect the weights of all its edges. We use a frequency map (or an array since weights are small) to count how many times each weight appears. The total number of edges on the path is also counted. To minimize the number of edge modifications, we should change all edge weights to match the most frequent weight on the path. If the most frequent weight appears maxFreq times and the total path length is pathLength, we need to change pathLength - maxFreq edges. This process is repeated for every single query.

1import java.util.*;2 3class Solution {4    public int[] minOperationsQueries(int n, int[][] edges, int[][] queries) {5        List<int[]>[] adj = new ArrayList[n];6        for (int i = 0; i < n; i++) {7            adj[i] = new ArrayList<>();8        }9        for (int[] edge : edges) {10            adj[edge[0]].add(new int[]{edge[1], edge[2]});11            adj[edge[1]].add(new int[]{edge[0], edge[2]});12        }13 14        int[] result = new int[queries.length];15        for (int i = 0; i < queries.length; i++) {16            int u = queries[i][0];17            int v = queries[i][1];18 19            if (u == v) {20                result[i] = 0;21                continue;22            }23 24            // Find path from u to v using BFS25            int[] parent = new int[n];26            int[] edgeWeightToParent = new int[n];27            Arrays.fill(parent, -1);28            Queue<Integer> q = new LinkedList<>();29            30            q.offer(u);31            parent[u] = u; // Mark u as visited and its own parent32 33            while (!q.isEmpty()) {34                int curr = q.poll();35                if (curr == v) {36                    break;37                }38                for (int[] neighbor : adj[curr]) {39                    int nextNode = neighbor[0];40                    int weight = neighbor[1];41                    if (parent[nextNode] == -1) {42                        parent[nextNode] = curr;43                        edgeWeightToParent[nextNode] = weight;44                        q.offer(nextNode);45                    }46                }47            }48 49            // Backtrack from v to u to get path weights50            Map<Integer, Integer> freq = new HashMap<>();51            int pathLength = 0;52            int curr = v;53            while (curr != u) {54                int w = edgeWeightToParent[curr];55                freq.put(w, freq.getOrDefault(w, 0) + 1);56                pathLength++;57                curr = parent[curr];58            }59 60            int maxFreq = 0;61            for (int count : freq.values()) {62                maxFreq = Math.max(maxFreq, count);63            }64 65            result[i] = pathLength - maxFreq;66        }67        return result;68    }69}

Complexity

Time

O(M * N), where M is the number of queries and N is the number of nodes. For each of the M queries, we perform a traversal (BFS/DFS) which takes O(N) time in a tree.

Space

O(N), where N is the number of nodes. This is for storing the adjacency list, the parent array, and the queue for BFS for each query.

Trade-offs

Pros

Simple to understand and implement.
Requires no complex data structures or algorithms beyond basic graph traversal.

Cons

Highly inefficient for a large number of queries.
The time complexity of O(M * N) is too slow for the given constraints and will likely result in a Time Limit Exceeded (TLE) error.

A much more efficient approach involves precomputation. The key observation is that the path between any two nodes a and b goes through their Lowest Common Ancestor (LCA). By precomputing information about paths from the root to every node, we can answer queries about paths between any two nodes quickly. Specifically, we can precompute the counts of each edge weight on the path from the root to every node. This turns the problem of finding frequencies on a path a-b into a constant number of lookups using the precomputed values for a, b, and lca(a, b).

Algorithm

Preprocessing: a. Build an adjacency list for the tree. b. Perform a single DFS traversal starting from the root (e.g., node 0). During the traversal, compute: i. depth[u]: The depth of each node u. ii. parent[u][0]: The direct parent of u (the 2^0-th ancestor). iii. counts[u][w]: An array of prefix counts, where counts[u][w] is the number of edges with weight w on the path from the root to u. c. Build a binary lifting table parent[u][j] which stores the 2^j-th ancestor of u. This can be done with dynamic programming: parent[u][j] = parent[parent[u][j-1]][j-1].
Query Processing: a. For each query (a, b): b. Find the Lowest Common Ancestor, l = lca(a, b), using the precomputed binary lifting table. This takes O(log N) time. c. The path from a to b is composed of the path from a to l and the path from b to l. The total number of edges on this path is pathLength = depth[a] + depth[b] - 2 * depth[l]. d. The frequency of any weight w on the path from a to b can be calculated using the precomputed prefix counts: freq(w) = counts[a][w] + counts[b][w] - 2 * counts[l][w]. e. Iterate through all possible weights (1 to 26) to find the maximum frequency maxFreq on the path. f. The result for the query is pathLength - maxFreq.

Walkthrough

This approach is divided into two phases: preprocessing and query answering.

Preprocessing: First, we traverse the tree using DFS from a root (say, node 0). This traversal allows us to compute several key pieces of information for each node u:

depth[u]: The distance from the root to u.
parent[u][0]: The immediate parent of u.
counts[u]: An array where counts[u][w] stores the total count of edges with weight w on the path from the root to u. This is a form of prefix sum on the tree.

Next, we use the parent[u][0] information to build a binary lifting table, parent[u][j], which allows us to find the 2^j-th ancestor of any node u in O(1) time. Building this table takes O(N log N) time.

Query Answering: For each query (a, b):

We first find their Lowest Common Ancestor, l = lca(a, b), in O(log N) time using the binary lifting table.
The total number of edges on the path a -> b is depth[a] + depth[b] - 2 * depth[l].
The count of a specific weight w on the path a -> b is the sum of counts on paths root -> a and root -> b, minus the counts on the overlapping path root -> l counted twice. So, freq(w) = counts[a][w] + counts[b][w] - 2 * counts[l][w].
We compute this for all possible weights (1 to 26) to find the maximum frequency maxFreq.
The final answer is pathLength - maxFreq.

This method processes each query very quickly after an initial setup cost.

1import java.util.*;2 3class Solution {4    private static final int MAX_WEIGHT = 27;5    private List<int[]>[] adj;6    private int[][] parent;7    private int[][] counts;8    private int[] depth;9    private int LOG_N;10 11    public int[] minOperationsQueries(int n, int[][] edges, int[][] queries) {12        adj = new ArrayList[n];13        for (int i = 0; i < n; i++) {14            adj[i] = new ArrayList<>();15        }16        for (int[] edge : edges) {17            adj[edge[0]].add(new int[]{edge[1], edge[2]});18            adj[edge[1]].add(new int[]{edge[0], edge[2]});19        }20 21        LOG_N = (int) (Math.log(n) / Math.log(2)) + 1;22        parent = new int[n][LOG_N];23        for (int i = 0; i < n; i++) {24            Arrays.fill(parent[i], -1);25        }26        counts = new int[n][MAX_WEIGHT];27        depth = new int[n];28 29        dfs(0, -1, 0);30 31        for (int j = 1; j < LOG_N; j++) {32            for (int i = 0; i < n; i++) {33                if (parent[i][j - 1] != -1) {34                    parent[i][j] = parent[parent[i][j - 1]][j - 1];35                }36            }37        }38 39        int[] result = new int[queries.length];40        for (int i = 0; i < queries.length; i++) {41            int u = queries[i][0];42            int v = queries[i][1];43            44            if (u == v) {45                result[i] = 0;46                continue;47            }48 49            int lcaNode = lca(u, v);50            int pathLength = depth[u] + depth[v] - 2 * depth[lcaNode];51            52            int maxFreq = 0;53            for (int w = 1; w < MAX_WEIGHT; w++) {54                int freqOnPath = counts[u][w] + counts[v][w] - 2 * counts[lcaNode][w];55                maxFreq = Math.max(maxFreq, freqOnPath);56            }57            58            result[i] = pathLength - maxFreq;59        }60        return result;61    }62 63    private void dfs(int u, int p, int d) {64        parent[u][0] = p;65        depth[u] = d;66        for (int[] neighbor : adj[u]) {67            int v = neighbor[0];68            int w = neighbor[1];69            if (v != p) {70                System.arraycopy(counts[u], 0, counts[v], 0, MAX_WEIGHT);71                counts[v][w]++;72                dfs(v, u, d + 1);73            }74        }75    }76 77    private int lca(int u, int v) {78        if (depth[u] < depth[v]) {79            int temp = u; u = v; v = temp;80        }81 82        for (int j = LOG_N - 1; j >= 0; j--) {83            if (depth[u] - (1 << j) >= depth[v]) {84                u = parent[u][j];85            }86        }87 88        if (u == v) return u;89 90        for (int j = LOG_N - 1; j >= 0; j--) {91            if (parent[u][j] != -1 && parent[u][j] != parent[v][j]) {92                u = parent[u][j];93                v = parent[v][j];94            }95        }96        return parent[u][0];97    }98}

Complexity

Time

O((N+M) * (log N + W_max)), where N is nodes, M is queries, and W_max is max weight. Preprocessing takes O(N * (log N + W_max)) for DFS and building tables. Each query takes O(log N + W_max) for LCA and frequency calculation.

Space

O(N * (log N + W_max)), where N is the number of nodes and W_max is the maximum weight value. This space is used for the adjacency list, depth array, binary lifting table `parent` (O(N log N)), and the prefix counts table `counts` (O(N * W_max)).

Trade-offs

Pros

Extremely efficient for a large number of queries.
The 'query' part of the time complexity is very low, making it suitable for online query processing.

Cons

More complex to implement, requiring knowledge of advanced tree algorithms like LCA and binary lifting.
Requires significant preprocessing time and space, which might be overkill for a very small number of queries.

Solutions

Solution

1class Solution {2public3  int[] minOperationsQueries(int n, int[][] edges, int[][] queries) {4    int m = 32 - Integer.numberOfLeadingZeros(n);5    List<int[]>[] g = new List[n];6    Arrays.setAll(g, i->new ArrayList<>());7    int[][] f = new int[n][m];8    int[] p = new int[n];9    int[][] cnt = new int[n][0];10    int[] depth = new int[n];11    for (var e : edges) {12      int u = e[0], v = e[1], w = e[2] - 1;13      g[u].add(new int[]{v, w});14      g[v].add(new int[]{u, w});15    }16    cnt[0] = new int[26];17    Deque<Integer> q = new ArrayDeque<>();18    q.offer(0);19    while (!q.isEmpty()) {20      int i = q.poll();21      f[i][0] = p[i];22      for (int j = 1; j < m; ++j) {23        f[i][j] = f[f[i][j - 1]][j - 1];24      }25      for (var nxt : g[i]) {26        int j = nxt[0], w = nxt[1];27        if (j != p[i]) {28          p[j] = i;29          cnt[j] = cnt[i].clone();30          cnt[j][w]++;31          depth[j] = depth[i] + 1;32          q.offer(j);33        }34      }35    }36    int k = queries.length;37    int[] ans = new int[k];38    for (int i = 0; i < k; ++i) {39      int u = queries[i][0], v = queries[i][1];40      int x = u, y = v;41      if (depth[x] < depth[y]) {42        int t = x;43        x = y;44        y = t;45      }46      for (int j = m - 1; j >= 0; --j) {47        if (depth[x] - depth[y] >= (1 << j)) {48          x = f[x][j];49        }50      }51      for (int j = m - 1; j >= 0; --j) {52        if (f[x][j] != f[y][j]) {53          x = f[x][j];54          y = f[y][j];55        }56      }57      if (x != y) {58        x = p[x];59      }60      int mx = 0;61      for (int j = 0; j < 26; ++j) {62        mx = Math.max(mx, cnt[u][j] + cnt[v][j] - 2 * cnt[x][j]);63      }64      ans[i] = depth[u] + depth[v] - 2 * depth[x] - mx;65    }66    return ans;67  }68}

1class Solution {2public:3  vector<int> minOperationsQueries(int n, vector<vector<int>> &edges,4                                   vector<vector<int>> &queries) {5    int m = 32 - __builtin_clz(n);6    vector<pair<int, int>> g[n];7    int f[n][m];8    int p[n];9    int cnt[n][26];10    int depth[n];11    memset(f, 0, sizeof(f));12    memset(cnt, 0, sizeof(cnt));13    memset(depth, 0, sizeof(depth));14    memset(p, 0, sizeof(p));15    for (auto &e : edges) {16      int u = e[0], v = e[1], w = e[2] - 1;17      g[u].emplace_back(v, w);18      g[v].emplace_back(u, w);19    }20    queue<int> q;21    q.push(0);22    while (!q.empty()) {23      int i = q.front();24      q.pop();25      f[i][0] = p[i];26      for (int j = 1; j < m; ++j) {27        f[i][j] = f[f[i][j - 1]][j - 1];28      }29      for (auto &[j, w] : g[i]) {30        if (j != p[i]) {31          p[j] = i;32          memcpy(cnt[j], cnt[i], sizeof(cnt[i]));33          cnt[j][w]++;34          depth[j] = depth[i] + 1;35          q.push(j);36        }37      }38    }39    vector<int> ans;40    for (auto &qq : queries) {41      int u = qq[0], v = qq[1];42      int x = u, y = v;43      if (depth[x] < depth[y]) {44        swap(x, y);45      }46      for (int j = m - 1; ~j; --j) {47        if (depth[x] - depth[y] >= (1 << j)) {48          x = f[x][j];49        }50      }51      for (int j = m - 1; ~j; --j) {52        if (f[x][j] != f[y][j]) {53          x = f[x][j];54          y = f[y][j];55        }56      }57      if (x != y) {58        x = p[x];59      }60      int mx = 0;61      for (int j = 0; j < 26; ++j) {62        mx = max(mx, cnt[u][j] + cnt[v][j] - 2 * cnt[x][j]);63      }64      ans.push_back(depth[u] + depth[v] - 2 * depth[x] - mx);65    }66    return ans;67  }68};

1class Solution:2    def minOperationsQueries(self, n: int, edges: List[List[int]], queries: List[List[int]]) -> List[int]: m = n . bit_length() g = [[] for _ in range(n)] f = [[0] * m for _ in range(n)] p = [0] * n cnt = [None] * n depth = [0] * n for u, v, w in edges: g[u]. append((v, w - 1)) g[v]. append((u, w - 1)) cnt[0] = [0] * 26 q = deque([0]) while q: i = q . popleft() f[i][0] = p[i] for j in range(1, m): f[i][j] = f[f[i][j - 1]][j - 1] for j, w in g[i]: if j != p[i]: p[j] = i cnt[j] = cnt[i][:] cnt[j][w] += 1 depth[j] = depth[i] + 1 q . append(j) ans = [] for u, v in queries: x, y = u, v if depth[x] < depth[y]: x, y = y, x for j in reversed(range(m)): if depth[x] - depth[y] >= (1 << j): x = f[x][j] for j in reversed(range(m)): if f[x][j] != f[y][j]: x, y = f[x][j], f[y][j] if x != y: x = p[x] mx = max(cnt[u][j] + cnt[v][j] - 2 * cnt[x][j] for j in range(26)) ans . append(depth[u] + depth[v] - 2 * depth[x] - mx) return ans

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem