leetcode 3559. Number of Ways to Assign Edge Weights II

• leetcode 3559. Number of Ways to Assign Edge Weights II
  • 1. 解题思路
  • 2. 代码实现

• 题目链接：3559. Number of Ways to Assign Edge Weights II

1. 解题思路

这一题是题目3558. Number of Ways to Assign Edge Weights I的进阶版本。

对于题目3558来说，其核心就是找出树的最大深度，然后就是一个数学题了，要使得这一条路径上的边的总和为奇数，那么其可选的方式为：
$$N=\sum _{i=1}^{i\le n,i=1,3,5,\cdots }{C}_n^i=2^{n-1}$$

因此，我们可以直接计算出其结果。

而对于本题，问题增加的难度在于需要快速地任意给到的两个点$u,v$，找出这两点间的连接线段的长度，而这个的话事实上和题目3553. Minimum Weighted Subgraph With the Required Paths II非常相近，事实上我们只需要分别求出这两个点$u,v$的深度以及他们的最小公共父节点$p$的深度，那么$u,v$间的距离$d$就可以快速通过下述公式计算得到：
$$d=h_u+h_v-2\times h_p$$

剩下的就是如何求出$u,v$的最小公共父节点了，而这个事实上就是一个经典算法了，这里就不具体展开了，后续有时间的话可能会系统整理一下这个算法然后发一篇博客作为备忘好了。

2. 代码实现

给出python代码实现如下：

MOD = 10**9+7

class LCA:
    def __init__(self, root, tree):
        self.max_level = 20  # 根据树的高度调整
        self.parent = defaultdict(lambda: [-1]*self.max_level)
        self.depth = {}
        self.preprocess(root, tree)

    def preprocess(self, root, tree):
        stack = [(root, -1, 0)]  # (node, parent, depth)
        while stack:
            node, par, d = stack.pop()
            self.depth[node] = d
            self.parent[node][0] = par
            for k in range(1, self.max_level):
                if self.parent[node][k-1] != -1:
                    self.parent[node][k] = self.parent[self.parent[node][k-1]][k-1]
            for child in tree[node]:
                if child != par:
                    stack.append((child, node, d+1))

    def query(self, u, v):
        if self.depth[u] < self.depth[v]:
            u, v = v, u
        # 对齐深度
        for k in range(self.max_level-1, -1, -1):
            if self.depth[u] - (1 << k) >= self.depth[v]:
                u = self.parent[u][k]
        if u == v:
            return u
        # 同步跳转
        for k in range(self.max_level-1, -1, -1):
            if self.parent[u][k] != -1 and self.parent[u][k] != self.parent[v][k]:
                u = self.parent[u][k]
                v = self.parent[v][k]
        return self.parent[u][0]
    
class Solution:
    def assignEdgeWeights(self, edges: List[List[int]], queries: List[List[int]]) -> List[int]:
        graph = defaultdict(list)
        for u, v in edges:
            graph[u].append(v)
            graph[v].append(u)
        
        depth = defaultdict(int)
        tree = defaultdict(list)
        def dfs(u, p):
            nonlocal depth, tree
            if u == 1:
                depth[u] = 0
            else:
                depth[u] = depth[p] + 1
            for v in graph[u]:
                if v == p:
                    continue
                tree[u].append(v)
                dfs(v, u)
            return
        
        dfs(1, 0)
        lca = LCA(1, tree)
        
        def query(u, v):
            p = lca.query(u, v)
            d = depth[u] + depth[v] - 2*depth[p]
            return pow(2, d-1, MOD) if d > 0 else 0
        
        return [query(u, v) for u, v in queries]

提交代码评测得到：耗时2773ms，占用内存137.3MB。