3559. Number of Ways to Assign Edge Weights II

import java.util.*;

class Solution {

    static final int MOD = 1_000_000_007;
    static final int LOG = 17;

    List<Integer>[] graph;
    int[][] up;
    int[] depth;

    public int[] assignEdgeWeights(int[][] edges, int[][] queries) {

        int n = edges.length + 1;

        graph = new ArrayList[n + 1];
        for (int i = 1; i <= n; i++) {
            graph[i] = new ArrayList<>();
        }

        for (int[] e : edges) {
            graph[e[0]].add(e[1]);
            graph[e[1]].add(e[0]);
        }

        up = new int[n + 1][LOG + 1];
        depth = new int[n + 1];

        dfs(1, 0);

        for (int j = 1; j <= LOG; j++) {
            for (int i = 1; i <= n; i++) {
                up[i][j] = up[up[i][j - 1]][j - 1];
            }
        }

        int maxDepth = n;
        long[] pow2 = new long[maxDepth + 1];
        pow2[0] = 1;

        for (int i = 1; i <= maxDepth; i++) {
            pow2[i] = (pow2[i - 1] * 2) % MOD;
        }

        int[] ans = new int[queries.length];

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

            int u = queries[i][0];
            int v = queries[i][1];

            int lca = getLCA(u, v);

            int len = depth[u] + depth[v] - 2 * depth[lca];

            if (len == 0) {
                ans[i] = 0;
            } else {
                ans[i] = (int) pow2[len - 1];
            }
        }

        return ans;
    }

    private void dfs(int node, int parent) {

        up[node][0] = parent;

        for (int nxt : graph[node]) {
            if (nxt == parent) continue;

            depth[nxt] = depth[node] + 1;
            dfs(nxt, node);
        }
    }

    private int getLCA(int u, int v) {

        if (depth[u] < depth[v]) {
            int temp = u;
            u = v;
            v = temp;
        }

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

        for (int j = LOG; j >= 0; j--) {
            if (((diff >> j) & 1) == 1) {
                u = up[u][j];
            }
        }

        if (u == v) return u;

        for (int j = LOG; j >= 0; j--) {
            if (up[u][j] != up[v][j]) {
                u = up[u][j];
                v = up[v][j];
            }
        }

        return up[u][0];
    }
}

class Solution:
    MOD = 1_000_000_007
    LOG = 17

    def assignEdgeWeights(self, edges, queries):

        n = len(edges) + 1

        graph = [[] for _ in range(n + 1)]

        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)

        up = [[0] * (self.LOG + 1) for _ in range(n + 1)]
        depth = [0] * (n + 1)

        def dfs(node, parent):
            up[node][0] = parent

            for nxt in graph[node]:
                if nxt == parent:
                    continue

                depth[nxt] = depth[node] + 1
                dfs(nxt, node)

        dfs(1, 0)

        for j in range(1, self.LOG + 1):
            for i in range(1, n + 1):
                up[i][j] = up[up[i][j - 1]][j - 1]

        pow2 = [1] * (n + 1)

        for i in range(1, n + 1):
            pow2[i] = (pow2[i - 1] * 2) % self.MOD

        def lca(u, v):

            if depth[u] < depth[v]:
                u, v = v, u

            diff = depth[u] - depth[v]

            for j in range(self.LOG, -1, -1):
                if (diff >> j) & 1:
                    u = up[u][j]

            if u == v:
                return u

            for j in range(self.LOG, -1, -1):
                if up[u][j] != up[v][j]:
                    u = up[u][j]
                    v = up[v][j]

            return up[u][0]

        ans = []

        for u, v in queries:

            p = lca(u, v)

            length = depth[u] + depth[v] - 2 * depth[p]

            ans.append(0 if length == 0 else pow2[length - 1])

        return ans

class Solution {
public:

    static constexpr int MOD = 1000000007;
    static constexpr int LOG = 17;

    vector<vector<int>> graph;
    vector<vector<int>> up;
    vector<int> depth;

    void dfs(int node,int parent){

        up[node][0]=parent;

        for(int nxt:graph[node]){

            if(nxt==parent) continue;

            depth[nxt]=depth[node]+1;

            dfs(nxt,node);
        }
    }

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

        int n=edges.size()+1;

        graph.assign(n+1,{});

        for(auto &e:edges){

            graph[e[0]].push_back(e[1]);
            graph[e[1]].push_back(e[0]);
        }

        up.assign(n+1,vector<int>(LOG+1));

        depth.assign(n+1,0);

        dfs(1,0);

        for(int j=1;j<=LOG;j++)
            for(int i=1;i<=n;i++)
                up[i][j]=up[up[i][j-1]][j-1];

        vector<long long> pow2(n+1);

        pow2[0]=1;

        for(int i=1;i<=n;i++)
            pow2[i]=pow2[i-1]*2%MOD;

        vector<int> ans;

        for(auto &q:queries){

            int u=q[0],v=q[1];

            if(depth[u]<depth[v]) swap(u,v);

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

            for(int j=LOG;j>=0;j--)
                if((diff>>j)&1)
                    u=up[u][j];

            int p;

            if(u==v) p=u;
            else{

                for(int j=LOG;j>=0;j--)
                    if(up[u][j]!=up[v][j]){
                        u=up[u][j];
                        v=up[v][j];
                    }

                p=up[u][0];
            }

            int len=depth[q[0]]+depth[q[1]]-2*depth[p];

            ans.push_back(len==0?0:pow2[len-1]);
        }

        return ans;
    }
};

var assignEdgeWeights = function(edges, queries) {

    const MOD = 1000000007n;
    const LOG = 17;

    const n = edges.length + 1;

    const graph = Array.from(
        {length:n+1},
        ()=>[]
    );

    for(const [u,v] of edges){

        graph[u].push(v);
        graph[v].push(u);
    }

    const up = Array.from(
        {length:n+1},
        ()=>Array(LOG+1).fill(0)
    );

    const depth = Array(n+1).fill(0);

    function dfs(node,parent){

        up[node][0]=parent;

        for(const nxt of graph[node]){

            if(nxt===parent) continue;

            depth[nxt]=depth[node]+1;

            dfs(nxt,node);
        }
    }

    dfs(1,0);

    for(let j=1;j<=LOG;j++)
        for(let i=1;i<=n;i++)
            up[i][j]=up[up[i][j-1]][j-1];

    const pow2=Array(n+1).fill(1n);

    for(let i=1;i<=n;i++)
        pow2[i]=(pow2[i-1]*2n)%MOD;

    function lca(u,v){

        if(depth[u]<depth[v]) [u,v]=[v,u];

        let diff=depth[u]-depth[v];

        for(let j=LOG;j>=0;j--)
            if((diff>>j)&1)
                u=up[u][j];

        if(u===v) return u;

        for(let j=LOG;j>=0;j--)
            if(up[u][j]!==up[v][j]){
                u=up[u][j];
                v=up[v][j];
            }

        return up[u][0];
    }

    const ans=[];

    for(const [u,v] of queries){

        const p=lca(u,v);

        const len=
        depth[u]+depth[v]-2*depth[p];

        ans.push(
            len===0
            ?0
            :Number(pow2[len-1])
        );
    }

    return ans;
};