3559. Number of Ways to Assign Edge Weights II

• Time: $O(n\log n + q\log n)$
• Space: $O(n + q)$

C++JavaPython

class Solution {
 public:
  vector<int> assignEdgeWeights(vector<vector<int>>& edges,
                                vector<vector<int>>& queries) {
    const int n = edges.size() + 1;
    vector<int> ans;
    vector<int> depth(n + 1);
    vector<vector<int>> graph(n + 1);
    vector<vector<int>> parent(kLog, vector<int>(n + 1, -1));

    for (const vector<int>& edge : edges) {
      const int u = edge[0];
      const int v = edge[1];
      graph[u].push_back(v);
      graph[v].push_back(u);
    }

    dfs(1, -1, graph, parent, depth);

    for (int k = 1; k < kLog; ++k)
      for (int v = 1; v <= n; ++v)
        if (parent[k - 1][v] != -1)
          parent[k][v] = parent[k - 1][parent[k - 1][v]];

    for (const vector<int>& query : queries) {
      const int u = query[0];
      const int v = query[1];
      if (u == v) {
        ans.push_back(0);
      } else {
        const int a = lca(u, v, parent, depth);
        const int d = depth[u] + depth[v] - 2 * depth[a];
        ans.push_back(modPow(2, d - 1));
      }
    }

    return ans;
  }

 private:
  static constexpr int kMod = 1'000'000'007;
  static constexpr int kLog = 17;  // since 2^17 > 1e5

  void dfs(int u, int p, const vector<vector<int>>& graph,
           vector<vector<int>>& parent, vector<int>& depth) {
    parent[0][u] = p;
    for (const int v : graph[u])
      if (v != p) {
        depth[v] = depth[u] + 1;
        dfs(v, u, graph, parent, depth);
      }
  }

  int lca(int u, int v, const vector<vector<int>>& parent,
          const vector<int>& depth) {
    if (depth[u] < depth[v])
      swap(u, v);

    for (int k = kLog - 1; k >= 0; --k)
      if (parent[k][u] != -1 && depth[parent[k][u]] >= depth[v])
        u = parent[k][u];

    if (u == v)
      return u;

    for (int k = kLog - 1; k >= 0; --k)
      if (parent[k][u] != -1 && parent[k][u] != parent[k][v]) {
        u = parent[k][u];
        v = parent[k][v];
      }

    return parent[0][u];
  }

  long modPow(long x, long n) {
    if (n == 0)
      return 1;
    if (n % 2 == 1)
      return x * modPow(x % kMod, (n - 1)) % kMod;
    return modPow(x * x % kMod, (n / 2)) % kMod;
  }
};

class Solution {
  public int[] assignEdgeWeights(int[][] edges, int[][] queries) {
    final int n = edges.length + 1;
    int[] ans = new int[queries.length];
    final int[] depth = new int[n + 1];
    final int[][] parent = new int[LOG][n + 1];
    List<Integer>[] graph = new List[n + 1];
    Arrays.setAll(graph, i -> new ArrayList<>());

    for (int[] edge : edges) {
      final int u = edge[0];
      final int v = edge[1];
      graph[u].add(v);
      graph[v].add(u);
    }

    dfs(1, -1, graph, parent, depth);

    for (int k = 1; k < LOG; ++k)
      for (int v = 1; v <= n; ++v)
        if (parent[k - 1][v] != -1)
          parent[k][v] = parent[k - 1][parent[k - 1][v]];

    for (int i = 0; i < queries.length; ++i) {
      final int u = queries[i][0];
      final int v = queries[i][1];
      if (u == v) {
        ans[i] = 0;
      } else {
        final int a = lca(u, v, parent, depth);
        final int d = depth[u] + depth[v] - 2 * depth[a];
        ans[i] = modPow(2, d - 1);
      }
    }

    return ans;
  }

  private static final int MOD = 1_000_000_007;
  private static final int LOG = 17; // since 2^17 > 1e5

  private void dfs(int u, int p, java.util.List<Integer>[] graph, int[][] parent, int[] depth) {
    parent[0][u] = p;
    for (int v : graph[u]) {
      if (v != p) {
        depth[v] = depth[u] + 1;
        dfs(v, u, graph, parent, depth);
      }
    }
  }

  private int lca(int u, int v, int[][] parent, int[] depth) {
    if (depth[u] < depth[v]) {
      final int temp = u;
      u = v;
      v = temp;
    }

    for (int k = LOG - 1; k >= 0; --k)
      if (parent[k][u] != -1 && depth[parent[k][u]] >= depth[v])
        u = parent[k][u];

    if (u == v)
      return u;

    for (int k = LOG - 1; k >= 0; --k)
      if (parent[k][u] != -1 && parent[k][u] != parent[k][v]) {
        u = parent[k][u];
        v = parent[k][v];
      }

    return parent[0][u];
  }

  private int modPow(long x, long n) {
    if (n == 0)
      return 1;
    if (n % 2 == 1)
      return (int) (x * modPow(x % MOD, (n - 1)) % MOD);
    return modPow(x * x % MOD, (n / 2)) % MOD;
  }
}

class Solution:
  def assignEdgeWeights(
      self,
      edges: list[list[int]],
      queries: list[list[int]]
  ) -> list[int]:
    MOD = 1_000_000_007
    LOG = 17  # since 2^17 > 1e5
    n = len(edges) + 1
    ans = []
    depth = [0] * (n + 1)
    graph = [[] for _ in range(n + 1)]
    parent = [[-1] * (n + 1) for _ in range(LOG)]

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

    def dfs(u: int, p: int) -> None:
      parent[0][u] = p
      for v in graph[u]:
        if v != p:
          depth[v] = depth[u] + 1
          dfs(v, u)

    dfs(1, -1)

    for k in range(1, LOG):
      for v in range(1, n + 1):
        if parent[k - 1][v] != -1:
          parent[k][v] = parent[k - 1][parent[k - 1][v]]

    def lca(u: int, v: int) -> int:
      if depth[u] < depth[v]:
        u, v = v, u

      for k in reversed(range(LOG)):
        if parent[k][u] != -1 and depth[parent[k][u]] >= depth[v]:
          u = parent[k][u]

      if u == v:
        return u

      for k in reversed(range(LOG)):
        if parent[k][u] != -1 and parent[k][u] != parent[k][v]:
          u = parent[k][u]
          v = parent[k][v]

      return parent[0][u]

    for u, v in queries:
      if u == v:
        ans.append(0)
      else:
        a = lca(u, v)
        d = depth[u] + depth[v] - 2 * depth[a]
        ans.append(pow(2, d - 1, MOD))

    return ans