【Leetcode】2846. 边权重均等查询

题目

题目链接
现有一棵由 n 个节点组成的无向树，节点按从 0 到 n - 1 编号。给你一个整数 n 和一个长度为 n - 1 的二维整数数组 edges ，其中 edges[i] = [ui, vi, wi] 表示树中存在一条位于节点 ui 和节点 vi 之间、权重为 wi 的边。

另给你一个长度为 m 的二维整数数组 queries ，其中 queries[i] = [ai, bi] 。对于每条查询，请你找出使从 ai 到 bi 路径上每条边的权重相等所需的 最小操作次数 。在一次操作中，你可以选择树上的任意一条边，并将其权重更改为任意值。

注意：

查询之间 相互独立 的，这意味着每条新的查询时，树都会回到 初始状态 。
从 ai 到 bi的路径是一个由 不同 节点组成的序列，从节点 ai 开始，到节点 bi 结束，且序列中相邻的两个节点在树中共享一条边。
返回一个长度为 m 的数组 answer ，其中 answer[i] 是第 i 条查询的答案。

示例 1：

输入：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]]
输出：[0,0,1,3]
解释：第 1 条查询，从节点 0 到节点 3 的路径中的所有边的权重都是 1 。因此，答案为 0 。
第 2 条查询，从节点 3 到节点 6 的路径中的所有边的权重都是 2 。因此，答案为 0 。
第 3 条查询，将边 [2,3] 的权重变更为 2 。在这次操作之后，从节点 2 到节点 6 的路径中的所有边的权重都是 2 。因此，答案为 1 。
第 4 条查询，将边 [0,1]、[1,2]、[2,3] 的权重变更为 2 。在这次操作之后，从节点 0 到节点 6 的路径中的所有边的权重都是 2 。因此，答案为 3 。
对于每条查询 queries[i] ，可以证明 answer[i] 是使从 ai 到 bi 的路径中的所有边的权重相等的最小操作次数。

示例 2：

输入：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]]
输出：[1,2,2,3]
解释：第 1 条查询，将边 [1,3] 的权重变更为 6 。在这次操作之后，从节点 4 到节点 6 的路径中的所有边的权重都是 6 。因此，答案为 1 。
第 2 条查询，将边 [0,3]、[3,1] 的权重变更为 6 。在这次操作之后，从节点 0 到节点 4 的路径中的所有边的权重都是 6 。因此，答案为 2 。
第 3 条查询，将边 [1,3]、[5,2] 的权重变更为 6 。在这次操作之后，从节点 6 到节点 5 的路径中的所有边的权重都是 6 。因此，答案为 2 。
第 4 条查询，将边 [0,7]、[0,3]、[1,3] 的权重变更为 6 。在这次操作之后，从节点 7 到节点 4 的路径中的所有边的权重都是 6 。因此，答案为 3 。
对于每条查询 queries[i] ，可以证明 answer[i] 是使从 ai 到 bi 的路径中的所有边的权重相等的最小操作次数。

提示：

• 1 <= n <= 10⁴
• edges.length == n - 1
• edges[i].length == 3
• 0 <= ui, vi < n
• 1 <= wi <= 26
• 生成的输入满足 edges 表示一棵有效的树
• 1 <= queries.length == m <= 2 * 10⁴
• queries[i].length == 2
• 0 <= ai, bi < n

思路

1. 将树的根节点设置为任意节点。
2. 定义一个二维数组 freq[node][weight]，用于保存从根到每个节点的路径上每条边权重的频率。
3. 从 a 到 b 的路径上边权重 w 的频率等于 freq[a][w] + freq[b][w] - freq[lca(a,b)][w] * 2，其中 lca(a,b) 是树中节点 a 和 b 的最低共同祖先。
4. 使用二进制提升算法来计算 lca(a,b)。

代码

class Solution {
public:
    vector minOperationsQueries(int n, vector>& edges, vector>& queries) {
        static const size_t L = 26;
        vector>> graph(n);
        for( auto &e : edges ) {
            graph[e[0]].emplace_back(e[1], e[2]-1);
            graph[e[1]].emplace_back(e[0], e[2]-1);
        }
        const int m = 32 - __builtin_clz(n);
        vector depth(n, 0);
        vector> ancestors(n, vector(m, -1));
        vector>> count(n, vector>(m));
        function dfs = [&](int i, int fa) {
            for( auto [j,x] : graph[i] ) {
                if( j == fa ) continue;
                depth[j] = depth[i] + 1;
                ancestors[j][0] = i;
                count[j][0][x] = 1;
                dfs(j, i);
            }
        };
        dfs(0, -1);

        for( int i = 1; i < m; ++i ) {
            for( int j = 0; j < n; ++j ) {
                int ac = ancestors[j][i-1]; 
                if( ac == -1 ) continue;
                ancestors[j][i] = ancestors[ac][i-1];
                for( int k = 0; k < L; ++k ) {
                    count[j][i][k] = count[j][i-1][k] + count[ac][i-1][k];
                }
            }
        }

        vector result;
        result.reserve(queries.size());
        for( auto &q : queries ) {
            int x = q[0], y = q[1];
            array mark{0};
            if( depth[x] > depth[y] ) swap(x,y);
            int total_depth = depth[x] + depth[y];
            int diff = depth[y]-depth[x];
            for( int c = 0; diff; diff>>= 1, ++c ) {
                if( diff & 1 ) {
                    for( int k = 0; k < L; ++k ) {
                        mark[k] += count[y][c][k];
                    }
                    y = ancestors[y][c];
                }
            }
            int same_ac = x;
            if( x != y ) {
                for( int c = m-1; c >= 0; --c ) {
                    int ax = ancestors[x][c];
                    int ay = ancestors[y][c];
                    if( ax != ay ) {
                        for( int k = 0; k < L; ++k ) {
                            mark[k] += count[x][c][k] + count[y][c][k];
                        }
                        x = ax;
                        y = ay;
                    }
                }
                for( int k = 0; k < L; ++k ) {
                    mark[k] += count[x][0][k] + count[y][0][k];
                }
                same_ac = ancestors[x][0];
            }
            int r = total_depth - depth[same_ac]*2 - *max_element(mark.begin(), mark.end());
            result.push_back(r);
        }
        return result;
    }
};

结果