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来说，其核心就是找出树的最大深度，然后就是一个数学题了，要使得这一条路径上的边的总和为奇数，那么其可选的方式为：
$\sum\limits_{i=1}^{i \leq 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$ 就可以快速通过下述公式计算得到：
$h_u + h_v - 2\times h_p$

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

2. 代码实现

给出python代码实现如下：

MOD = 10**9+7class 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] = dself.parent[node][0] = parfor 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, treeif u == 1:depth[u] = 0else:depth[u] = depth[p] + 1for v in graph[u]:if v == p:continuetree[u].append(v)dfs(v, u)returndfs(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 0return [query(u, v) for u, v in queries]