Leetcode每日一题 —— 3559. 给边赋权值的方案数 II

力扣这个月真上强度了…这题也是很综合，要写对难度真的还是不小的。顺便还学习到了求 LCA（最近公共祖先）的倍增法。

思路

看上去和昨天那道题还有点像，找的其实依旧是两个节点间的路径长度，但不同的是蹦出来一个 queries，而且其规模可能很大，线性级的 BFS 是难以接受的。

也就是说本题的难点主要是，我们需要在 O(\log n) 级别的复杂度下找到无向树任意两点间的路径长度。

看了提示才知道能用 LCA 算法，而进一步了解得知 O(\log n) 复杂度的求 LCA 算法有一种倍增法 (Binary Lifting)。

最终还得写个快速幂。

为什么能用 LCA

找到树中两个节点 u, v 的最近公共祖先 a 后，【从根节点到 u 的距离】和【从根节点到 v 的距离】都包含了【从根节点到 a 的距离】。把根节点到 u 和 v 分别的距离之和减去两份的【从根节点到 a 的距离】，得到的就是 u 到 v 的路径长度了。

至于无向树中的根节点，可以随意选择一个。

倍增法主要思想

以往通常用的递归 LCA 解法是线性时间复杂度的，自底向上跳跃直至收束到某个祖先，每次实际只跳了一步。

而倍增法则是把跳的步数按二进制位分解了，比如跳 14 步，可以分解成 14 (0b1110) = 8 + 4 + 2，也就是可以先跳 8 步再跳 4 步，最后跳 2 步，以此大大减少跳跃次数。

• 因此需要预先生成每个节点处向上跳跃 2^j 步所能到达的节点，用的是 boosts[node][j] 数组，可以看代码注释。


• 当然也需要用深度 deps 数组去记录每个节点的深度。

首先查询的两个节点 u, v 深度可能不同，我们可以先从较深的节点向上跳跃（根据分解出来的步数，先跳大步），使得二个指针处于相同深度，接着二者再往上一起跳跃：

1. 让较深的节点指针先跳：先计算两个节点的深度差 diff，把 diff 按二进制位分解来进行跳跃。这里先跳大步、再跳小步和先跳小步、再跳大步都可以（比如 14 步，可以 8->4->2，也可以 2->4->8）。


2. 二者从相同深度一起向上跳跃：必须先跳大步再跳小步，如果跳了 2^j 步发现二者还没相遇，则可以放心跳；否则则要减小步数再试。

还有一个问题，从大步往小步来试，那最大的步数 2^{k-1} 的 k 可能是多少呢？按最差的情况来看，整个树首尾串联成链表，最大的 k=\lfloor{\log_2{n}}\rfloor+1。

代码

注释尽量写详细了，说不定咱几天后又忘记解法了…

class Solution {

public:

    vector<int> assignEdgeWeights(vector<vector<int>>& edges, vector<vector<int>>& queries) {

        // 多了一个 queries，要计算指定两个节点的分配方式数量

        // 树的性质决定 u_i 到 v_i 间只会有一条路径

        // 关键就是要想办法快速求出两个节点之前的路径长度

        int n=edges.size()+1; // n 个节点

        // 先建成无向树

        vector<vector<int>> adjList(n);

        for(auto& e:edges){

            // 转换为 0...n-1 编号

            adjList[e[0]-1].emplace_back(e[1]-1);

            adjList[e[1]-1].emplace_back(e[0]-1);

        }

        // 为了方便处理，随便选一个节点作为根

        int root=rand()%n;

        // 这里可以用到 LCA (倍增法)

        // 计算倍增 2^k 的最大 k 值

        int k=1;

        while((1<<k)<=n){

            k++;

        }

        // 先初始化查询所需的数组

        vector<bool> visited(n, false); // 每个节点是否被访问

        vector<int> deps(n); // 每个节点的深度

        vector<vector<int>> boosts(n,vector<int>(k)); // 倍增表

        // 预处理树

        visited[root]=true; // root 的父节点标记为已经访问

        deps[root]=0; // root 深度显然为 0

        boosts[root][0]=root; // root 的父节点是自己，boosts[node][j] 表示 node 向上跳 2^j 次到达的祖先节点

        // BFS 初始化 deps 和 boosts 数组

        queue<int> q;

        q.emplace(root);

        while(!q.empty()){

            int curr=q.front();

            q.pop();

            // 这里 boosts[curr][0]，即 curr 的父节点已经设置

            // 往后推直至 boosts[curr][k-1]

            for(int j=1;j<=k-1;j++){

                // curr 往上跳 2^j 步

                // 相当于先跳 2^(j-1) 步 (到达 boosts[curr][j-1])，再跳 2^(j-1) 步

                // 因此 boosts[curr][j]=boosts[ boosts[curr][j-1] ][j-1]

                boosts[curr][j]=boosts[boosts[curr][j-1]][j-1];

            }

            // 扫描邻居

            for(int node:adjList[curr]){

                if(visited[node]){

                    // 已经访问过就 pass，避免回头

                    continue;

                }

                visited[node]=true;

                // 邻居的深度 +1

                deps[node]=deps[curr]+1;

                // 因为不走回头路，curr 就相当于邻居的父节点（向上跳 2^0 次）

                boosts[node][0]=curr;

                q.emplace(node);

            }

        }

        // 快速找到 u 和 v 的 LCA 的方法

        auto lca=[&](int u,int v)->int{

            // u, v 深度可能不同，要先让更深的跳到相同高度

            // 这里为了方便处理，让 u 是更深的那一个

            if(deps[u]<deps[v]){

                swap(u,v);

            }

            // 看看 u 要跳多少步才能到 v

            int diff=deps[u]-deps[v];

            // 接下来就是倍增法的精髓了

            // 不是让 u 一步一步跳完 steps

            // 而是让 steps 按二进制位分解

            // 比如 14 = 8 + 4 + 2 = 2^3 + 2^2 + 2^1

            // 这里从大步还是小步开始跳都可以，只用跳 3 次，而不是 14 次

            for(int j=k-1;j>=0;j--){

                if((diff&(1<<j))>0){

                    // j 这个二进制位有一个 1

                    // 通过 boosts 快速取出 u 向上跳 2^j 步到达的位置

                    u=boosts[u][j];

                }

            }

            // u 跳完后发现和 v 重合了，那 v 就是 LCA

            if(u==v){

                return u;

            }

            // 还没有重合，咱俩接着一起跳

            for(int j=k-1;j>=0;j--){

                // 这里必须先跳大步再跳小步

                // 如果 u 和 v 都跳了 2^j 还没有重合，则可以放心往上跳

                // 如果重合了，可能有这种情况：  

                //        1

                //        |

                //        2

                //       / \

                //      3   4

                // 从 3 和 4 往上跳大步会先跳到 1，但这不是 LCA

                // 因此重合的时候就先不跳，而是试着缩小步数

                if(boosts[u][j]!=boosts[v][j]){

                    u=boosts[u][j];

                    v=boosts[v][j];

                }

            }

            // 这样跳完后，保证 u 和 v 的父节点就是 LCA

            return boosts[u][0];

        };

        // 总算！我们能快速拿到两个节点的 LCA 后就有办法算两个节点的距离了！！

        // dist(u, v) = deps[u] + deps[v] - 2*deps[LCA(u, v)]

        // LCA 的深度是 u 和 v 共享的深度部分，去掉两份就得到 u 和 v 的距离了

        vector<int> res(queries.size());

        // 孩子们，别忘了快速幂！

        auto qPow=[&](int base,int exp)->int{

            long long res=1;

            long long b=base;

            while(exp>0){

                if((exp&1)==1){

                    res=res*b%(long long)(1e9+7);

                }

                b=(b*b)%(long long)(1e9+7);

                exp>>=1;

            }

            return res;

        };

        for(int i=0;i<queries.size();i++){

            int dist=deps[queries[i][0]-1]+deps[queries[i][1]-1]-deps[lca(queries[i][0]-1,queries[i][1]-1)]*2;

            // 注意有坑！查询的 u 和 v 可能相等！

            if(dist==0){

                res[i]=0;

            }else{

                res[i]=qPow(2,dist-1);

            }

        }

        return res;

    }

};