Minimum Edge Weight Equilibrium Queries in a Tree

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)$：

class Solution {
public:
    vector<int> minOperationsQueries(int n, vector<vector<int>>& edges, vector<vector<int>>& queries) {
        vector<vector<vector<int>>> g(n + 1);
        for (auto &p : edges) {
            p[0]++, p[1]++;
            g[p[0]].push_back({p[1], p[2]});
            g[p[1]].push_back({p[0], p[2]});
        }
        vector<vector<int>> fa(n + 1, vector<int>(14));
        vector<int> dep(n + 1, 0x3f3f3f3f);
        dep[0] = 0, dep[1] = 1;
        queue<int> q({1});
        while (!q.empty()) {
            int t = q.front();
            q.pop();
            for (auto &p : g[t]) {
                if (dep[p[0]] > dep[t] + 1) {
                    dep[p[0]] = dep[t] + 1;
                    q.push(p[0]);
                    fa[p[0]][0] = t;
                    for (int i = 1; i <= 13; i++) {
                        fa[p[0]][i] = fa[fa[p[0]][i - 1]][i - 1];
                    }
                }
            }
        }
        vector<vector<int>> s(27, vector<int>(n + 1));
        for (int w = 1; w <= 26; w++) {
            function<void(int, int)> dfs = [&](int u, int pre) {
                for (auto &p : g[u]) {
                    if (p[0] != pre) {
                        s[w][p[0]] = s[w][u] + (p[1] == w);
                        dfs(p[0], u);
                    }
                }
            };
            dfs(1, -1);
        }
        vector<int> ans;
        function<int(int, int)> lca = [&](int a, int b) {
            if (dep[a] < dep[b]) swap(a, b);
            for (int i = 13; i >= 0; i--) {
                if (dep[fa[a][i]] >= dep[b]) a = fa[a][i];
            }
            if (a == b) return a;
            for (int i = 13; i >= 0; i--) {
                if (fa[a][i] != fa[b][i]) a = fa[a][i], b = fa[b][i];
            }
            return fa[a][0];
        };
        for (auto &p : queries) {
            p[0]++, p[1]++;
            int x = lca(p[0], p[1]), t = n;
            for (int w = 1; w <= 26; w++) {
                t = min(t, dep[p[0]] + dep[p[1]] - 2 * dep[x] - (s[w][p[0]] + s[w][p[1]] - 2 * s[w][x]));
            }
            ans.push_back(t);
        }
        return ans;
    }
};