M3558.给边赋权值的方案数 I

bfs, math, https://leetcode.cn/problems/number-of-ways-to-assign-edge-weights-i/description/

给你一棵 n 个节点的无向树，节点从 1 到 n 编号，树以节点 1 为根。树由一个长度为 n - 1 的二维整数数组 edges 表示，其中 edges[i] = [ui, vi] 表示在节点 ui 和 vi 之间有一条边。

一开始，所有边的权重为 0。你可以将每条边的权重设为 1 或 2。

两个节点 u 和 v 之间路径的 代价 是连接它们路径上所有边的权重之和。

选择任意一个 深度最大 的节点 x。返回从节点 1 到 x 的路径中，边权重之和为 奇数 的赋值方式数量。

由于答案可能很大，返回它对 10^9 + 7 取模的结果。

注意： 忽略从节点 1 到节点 x 的路径外的所有边。

示例 1：

输入： edges = [[1,2]]

输出： 1

解释：

• 从节点 1 到节点 2 的路径有一条边（1 → 2）。
• 将该边赋权为 1 会使代价为奇数，赋权为 2 则为偶数。因此，合法的赋值方式有 1 种。

示例 2：

输入： edges = [[1,2],[1,3],[3,4],[3,5]]

输出： 2

解释：

• 最大深度为 2，节点 4 和节点 5 都在该深度，可以选择任意一个。
• 例如，从节点 1 到节点 4 的路径包括两条边（1 → 3 和 3 → 4）。
• 将两条边赋权为 (1,2) 或 (2,1) 会使代价为奇数，因此合法赋值方式有 2 种。

提示：

• 2 <= n <= 10^5

• edges.length == n - 1

• edges[i] == [ui, vi]

• 1 <= ui, vi <= n

• edges 表示一棵合法的树。

这道题结合了图论（树的遍历）与组合数学（二项式定理）。以下是对本题的详细解读、数学原理分析以及算法实现。

一、 题目核心解读

结合题目的定义与限制，可以将问题简化为以下几个步骤：

1. 确定目标路径的长度：

   • 题目要求选择任意一个最大深度的节点 $x$。
   • 深度定义为从根节点 1 到该节点的路径所包含的边数。
   • 设树的最大深度为 $d$。由于“忽略路径外的所有边”，实际上只需要关注这条长度为 $d$ 的路径，路径上有且仅有 $d$ 条边。

2. 边权赋值的奇偶性约束：

   • 每条边可以赋值为 1 或 2。
   • 路径的代价为所有边权之和。要使这个和为奇数。

二、 数学原理推导

设路径上共有 $d$ 条边。分配给这些边的值中，设共有 $k$ 条边被赋予了权重 1，剩下的 $d-k$ 条边被赋予了权重 2。

此时，路径的总权重之和为：

$$\text{Sum}=1\times k+2\times (d-k)=2d-k$$

要使 $\text{Sum}$ 为奇数，由于 $2d$ 显然是偶数，因此 $k$（权重为 1 的边数）必须为奇数。

这意味着，需要在 $d$ 条边中选择奇数个位置赋值为 1，其余位置赋值为 2。根据组合数，总方案数 $S$ 为：

$$S=\sum _{k\text{ 为奇数}}(\genfrac{}{}{0ex}{}{d}{k})=(\genfrac{}{}{0ex}{}{d}{1})+(\genfrac{}{}{0ex}{}{d}{3})+(\genfrac{}{}{0ex}{}{d}{5})+\dots$$

根据二项式定理的性质，在一个大小为 $d$ ($d\ge 1$) 的集合中，选择奇数个元素的子集个数与选择偶数个元素的子集个数相等，它们都等于总子集数（即 $2^d$）的一半：

$$S=2^{d-1}$$

因此，只要求出树的最大深度 $d$，最终的答案就是 $2^{d-1}(\mathrm{mod}{10}^9+7)$。因为 $n\ge 2$，所以 $d\ge 1$，公式始终有效。

三、 算法设计与实现

整个算法可以分为两步：

1. 求最大深度：使用广度优先搜索（BFS）或深度优先搜索（DFS）遍历整棵树，求出从根节点 1 出发所能到达的最大深度 $d$。
2. 计算结果：使用快速幂求出 $2^{d-1}(\mathrm{mod}{10}^9+7)$。

Python3 代码实现

from collections import deque
from typing import List

class Solution:
    def assignEdgeWeights(self, edges: List[List[int]]) -> int:
        n = len(edges) + 1
        # 1. 构建邻接表
        adj = [[] for _ in range(n + 1)]
        for u, v in edges:
            adj[u].append(v)
            adj[v].append(u)
        
        # 2. BFS 寻找从节点 1 出发的最深距离 (最大深度)
        max_depth = 0
        queue = deque([(1, 0)])  # 存储队列元素: (当前节点, 深度)
        visited = [False] * (n + 1)
        visited[1] = True
        
        while queue:
            u, depth = queue.popleft()
            if depth > max_depth:
                max_depth = depth
            for v in adj[u]:
                if not visited[v]:
                    visited[v] = True
                    queue.append((v, depth + 1))
        
        # 3. 计算结果
        MOD = 10**9 + 7
        return pow(2, max_depth - 1, MOD)

四、 复杂度分析

• 时间复杂度：$\mathcal{O}(n)$

  • 建图需要遍历所有的边，耗时 $\mathcal{O}(n)$。
  • BFS 遍历中，每个节点和每条边最多被访问常数次，耗时 $\mathcal{O}(n)$。
  • 快速幂计算 $2^{d-1}(\mathrm{mod}{10}^9+7)$ 的时间复杂度为 $\mathcal{O}(\log d)$。
  • 整体时间复杂度为 $\mathcal{O}(n)$，在 $n={10}^5$ 的数据规模下可以快速运行完毕。

• 空间复杂度：$\mathcal{O}(n)$

  • 邻接表需要 $\mathcal{O}(n)$ 的空间。
  • 队列和 visited 数组需要 $\mathcal{O}(n)$ 的空间。
  • 整体辅助空间复杂度为 $\mathcal{O}(n)$。