Minimum Edge Weight Equilibrium Queries in a Tree

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

解题思路

　　周赛很顺利做出这题，上大分，写篇题解纪念一下。

　　首先看到边的权值最大才$26$，因此很容易想到对于每个询问直接暴力枚举把边都变成哪个权值，然后再取最小的操作次数。关键是要知道两个节点$(u, v)$所构成的简单路径中每种权值的边共有多少。

　　假定树的根节点是$1$号点，直接暴力枚举所有可能的边权$w$，分别统计从$1$号节点出发到其他点的路径中有多少条权值恰好为$w$的边。令$s_{w,i}$表示从节点$1$到节点$i$的路径中权值为$w$的边的数量，因此如果要将$(u,v)$路径中的所有边的权值变成$w$，那么最小需要的操作次数就是$d_u + d_v - 2 \cdot d_p - (s_{w,u} + s_{w,v} - 2 \cdot s_{w,p})$，其中$p = \operatorname{lca}(u,v)$，$d_u$表示从节点$1$到节点$i$的路径长度。

　　所以对于每个询问$(u, v)$直接暴力枚举所有可能的$w$，然后取上式的最小值即可，即$$\min\limits_{1 \leq w \leq 26}\left\{d_u + d_v - 2 \cdot d_p - (s_{w,u} + s_{w,v} - 2 \cdot s_{w,p})\right\}$$

　　AC代码如下，时间复杂度为$O\left( n \log{n} + W \cdot n + q \cdot (\log{n} + W) \right)$：

 1 class Solution {
 2 public:
 3     vector<int> minOperationsQueries(int n, vector<vector<int>>& edges, vector<vector<int>>& queries) {
 4         vector<vector<vector<int>>> g(n + 1);
 5         for (auto &p : edges) {
 6             p[0]++, p[1]++;
 7             g[p[0]].push_back({p[1], p[2]});
 8             g[p[1]].push_back({p[0], p[2]});
 9         }
10         vector<vector<int>> fa(n + 1, vector<int>(14));
11         vector<int> dep(n + 1, 0x3f3f3f3f);
12         dep[0] = 0, dep[1] = 1;
13         queue<int> q({1});
14         while (!q.empty()) {
15             int t = q.front();
16             q.pop();
17             for (auto &p : g[t]) {
18                 if (dep[p[0]] > dep[t] + 1) {
19                     dep[p[0]] = dep[t] + 1;
20                     q.push(p[0]);
21                     fa[p[0]][0] = t;
22                     for (int i = 1; i <= 13; i++) {
23                         fa[p[0]][i] = fa[fa[p[0]][i - 1]][i - 1];
24                     }
25                 }
26             }
27         }
28         vector<vector<int>> s(27, vector<int>(n + 1));
29         for (int w = 1; w <= 26; w++) {
30             function<void(int, int)> dfs = [&](int u, int pre) {
31                 for (auto &p : g[u]) {
32                     if (p[0] != pre) {
33                         s[w][p[0]] = s[w][u] + (p[1] == w);
34                         dfs(p[0], u);
35                     }
36                 }
37             };
38             dfs(1, -1);
39         }
40         vector<int> ans;
41         function<int(int, int)> lca = [&](int a, int b) {
42             if (dep[a] < dep[b]) swap(a, b);
43             for (int i = 13; i >= 0; i--) {
44                 if (dep[fa[a][i]] >= dep[b]) a = fa[a][i];
45             }
46             if (a == b) return a;
47             for (int i = 13; i >= 0; i--) {
48                 if (fa[a][i] != fa[b][i]) a = fa[a][i], b = fa[b][i];
49             }
50             return fa[a][0];
51         };
52         for (auto &p : queries) {
53             p[0]++, p[1]++;
54             int x = lca(p[0], p[1]), t = n;
55             for (int w = 1; w <= 26; w++) {
56                 t = min(t, dep[p[0]] + dep[p[1]] - 2 * dep[x] - (s[w][p[0]] + s[w][p[1]] - 2 * s[w][x]));
57             }
58             ans.push_back(t);
59         }
60         return ans;
61     }
62 };