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 为根的有向树，用 BFS 或 DFS 计算每个节点到根的深度，取最大值记为 D。
计算方案数 只考虑从 1 到深度为 D 的某个节点的这条路径上的 D 条边，每条边权重只能是 1（奇）或 2（偶）。我们要统计总和为奇数的方案数。
令 O(D)O(D) 为长度为 DD 的序列中和为奇数的方案数，则有递推：
$O (D) = E (D - 1) \times (# 选 1) + O (D - 1) \times (# 选 2) = [2^{D - 1} - O (D - 1)] \times 1 + O (D - 1) \times 1 = 2^{D - 1}$ .
因此答案就是
$O (D) = 2^{D - 1} \mod (10^{9} + 7)$ .

python

from typing import List

class Solution:
    def assignEdgeWeights(self, edges: List[List[int]]) -> int:
        MOD = 10**9 + 7
        n = len(edges) + 1

        # 构建无向图
        g = [[] for _ in range(n+1)]
        for u, v in edges:
            g[u].append(v)
            g[v].append(u)

        # BFS 求各节点深度
        from collections import deque
        q = deque([1])
        depth = [0] * (n+1)
        seen = [False] * (n+1)
        seen[1] = True

        maxd = 0
        while q:
            u = q.popleft()
            for v in g[u]:
                if not seen[v]:
                    seen[v] = True
                    depth[v] = depth[u] + 1
                    maxd = max(maxd, depth[v])
                    q.append(v)

        return pow(2, maxd-1, MOD)

if __name__ == "__main__":
    sol = Solution()
    print(sol.assignEdgeWeights([[1,2]]))
    print(sol.assignEdgeWeights([[1,2],[1,3],[3,4],[3,5]]))

时间复杂度

构图 O(n)
BFS/DFS 求深度 O(n)
快速幂 O(logn) 总体 O(n)。

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

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