From 5e127269bf7762942a09c207c4eb24c891c66c3d Mon Sep 17 00:00:00 2001 From: Libin YANG Date: Thu, 11 Jun 2026 08:22:11 +0800 Subject: [PATCH] feat: add solutions to lc problem: No.3558 --- .../README.md | 185 +++++++++++++++++- .../README_EN.md | 185 +++++++++++++++++- .../Solution.cpp | 39 ++++ .../Solution.go | 37 ++++ .../Solution.java | 40 ++++ .../Solution.py | 16 ++ .../Solution.ts | 38 ++++ 7 files changed, 532 insertions(+), 8 deletions(-) create mode 100644 solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.cpp create mode 100644 solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.go create mode 100644 solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.java create mode 100644 solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.py create mode 100644 solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.ts diff --git a/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/README.md b/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/README.md index a40099f437bcb..289a3027a4566 100644 --- a/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/README.md +++ b/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/README.md @@ -88,32 +88,209 @@ tags: -### 方法一 +### 方法一：DFS + 数学 + +我们先通过边构建出 $g$，其中 $g[u]$ 表示 $u$ 的所有邻接点。 + +接下来，我们使用一个函数 $\textit{dfs}$ 来求出树的深度 $d$，那么从 $d$ 中选出奇数个数的方案，就是答案。根据定理，从 $d$ 中选出奇数个数的方案数为 $2^{d-1}$，我们可以用快速幂求出答案。 + +时间复杂度 $O(n)$，空间复杂度 $O(n)$。其中 $n$ 是树的节点数。 #### Python3 ```python - +class Solution: + def assignEdgeWeights(self, edges: List[List[int]]) -> int: + def dfs(i: int, fa: int = 0) -> int: + res = 0 + for j in g[i]: + if j != fa: + res = max(res, dfs(j, i) + 1) + return res + + n = len(edges) + 1 + g = [[] for _ in range(n + 1)] + for u, v in edges: + g[u].append(v) + g[v].append(u) + d = dfs(1) + return pow(2, d - 1, 10**9 + 7) ``` #### Java ```java - +class Solution { + private List[] g; + + public int assignEdgeWeights(int[][] edges) { + int n = edges.length + 1; + g = new List[n + 1]; + Arrays.setAll(g, k -> new ArrayList<>()); + + for (var e : edges) { + int u = e[0]; + int v = e[1]; + g[u].add(v); + g[v].add(u); + } + + return (int) pow(2, dfs(1, 0) - 1, 1_000_000_007); + } + + private int dfs(int i, int fa) { + int res = 0; + for (int j : g[i]) { + if (j != fa) { + res = Math.max(res, dfs(j, i) + 1); + } + } + return res; + } + + private long pow(long a, int n, int mod) { + long res = 1; + while (n > 0) { + if ((n & 1) != 0) { + res = res * a % mod; + } + a = a * a % mod; + n >>= 1; + } + return res; + } +} ``` #### C++ ```cpp - +class Solution { +public: + int assignEdgeWeights(vector>& edges) { + int n = edges.size() + 1; + vector> g(n + 1); + + for (auto& e : edges) { + int u = e[0]; + int v = e[1]; + g[u].push_back(v); + g[v].push_back(u); + } + + auto dfs = [&](this auto&& dfs, int i, int fa) -> int { + int res = 0; + for (int j : g[i]) { + if (j != fa) { + res = max(res, dfs(j, i) + 1); + } + } + return res; + }; + + return pow(2, dfs(1, 0) - 1, 1000000007); + } + +private: + long long pow(long long a, int n, int mod) { + long long res = 1; + while (n > 0) { + if (n & 1) { + res = res * a % mod; + } + a = a * a % mod; + n >>= 1; + } + return res; + } +}; ``` #### Go ```go +func assignEdgeWeights(edges [][]int) int { + const mod = 1_000_000_007 + + n := len(edges) + 1 + g := make([][]int, n+1) + + for _, e := range edges { + u, v := e[0], e[1] + g[u] = append(g[u], v) + g[v] = append(g[v], u) + } + + var dfs func(int, int) int + dfs = func(i, fa int) int { + res := 0 + for _, j := range g[i] { + if j != fa { + res = max(res, dfs(j, i)+1) + } + } + return res + } + + return pow(2, dfs(1, 0)-1, mod) +} + +func pow(a, n, mod int) int { + res := 1 + for n > 0 { + if n&1 > 0 { + res = res * a % mod + } + a = a * a % mod + n >>= 1 + } + return res +} +``` +#### TypeScript + +```ts +function assignEdgeWeights(edges: number[][]): number { + const mod = 1_000_000_007; + const n = edges.length + 1; + const g: number[][] = Array.from({ length: n + 1 }, () => []); + + for (const [u, v] of edges) { + g[u].push(v); + g[v].push(u); + } + + const dfs = (i: number, fa: number): number => { + let res = 0; + for (const j of g[i]) { + if (j !== fa) { + res = Math.max(res, dfs(j, i) + 1); + } + } + return res; + }; + + const pow = (a: number, n: number, mod: number): number => { + let res = 1n; + let x = BigInt(a); + const m = BigInt(mod); + + while (n > 0) { + if (n & 1) { + res = (res * x) % m; + } + x = (x * x) % m; + n >>= 1; + } + + return Number(res); + }; + + return pow(2, dfs(1, 0) - 1, mod); +} ``` diff --git a/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/README_EN.md b/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/README_EN.md index 9d87b1191c2a4..fcd55ff84d1a9 100644 --- a/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/README_EN.md +++ b/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/README_EN.md @@ -85,32 +85,209 @@ tags: -### Solution 1 +### Solution 1: DFS + Mathematics + +First, we build an adjacency list $g$ from the edges, where $g[u]$ contains all neighbors of node $u$. + +Next, we use a function $\textit{dfs}$ to compute the depth $d$ of the tree. The answer is the number of ways to choose an odd number of elements from $d$. According to a well-known combinatorial identity, the number of ways to choose an odd number of elements from a set of size $d$ is $2^{d-1}$. Therefore, we can compute the answer using fast exponentiation. + +The time complexity is $O(n)$, and the space complexity is $O(n)$, where $n$ is the number of nodes in the tree. #### Python3 ```python - +class Solution: + def assignEdgeWeights(self, edges: List[List[int]]) -> int: + def dfs(i: int, fa: int = 0) -> int: + res = 0 + for j in g[i]: + if j != fa: + res = max(res, dfs(j, i) + 1) + return res + + n = len(edges) + 1 + g = [[] for _ in range(n + 1)] + for u, v in edges: + g[u].append(v) + g[v].append(u) + d = dfs(1) + return pow(2, d - 1, 10**9 + 7) ``` #### Java ```java - +class Solution { + private List[] g; + + public int assignEdgeWeights(int[][] edges) { + int n = edges.length + 1; + g = new List[n + 1]; + Arrays.setAll(g, k -> new ArrayList<>()); + + for (var e : edges) { + int u = e[0]; + int v = e[1]; + g[u].add(v); + g[v].add(u); + } + + return (int) pow(2, dfs(1, 0) - 1, 1_000_000_007); + } + + private int dfs(int i, int fa) { + int res = 0; + for (int j : g[i]) { + if (j != fa) { + res = Math.max(res, dfs(j, i) + 1); + } + } + return res; + } + + private long pow(long a, int n, int mod) { + long res = 1; + while (n > 0) { + if ((n & 1) != 0) { + res = res * a % mod; + } + a = a * a % mod; + n >>= 1; + } + return res; + } +} ``` #### C++ ```cpp - +class Solution { +public: + int assignEdgeWeights(vector>& edges) { + int n = edges.size() + 1; + vector> g(n + 1); + + for (auto& e : edges) { + int u = e[0]; + int v = e[1]; + g[u].push_back(v); + g[v].push_back(u); + } + + auto dfs = [&](this auto&& dfs, int i, int fa) -> int { + int res = 0; + for (int j : g[i]) { + if (j != fa) { + res = max(res, dfs(j, i) + 1); + } + } + return res; + }; + + return pow(2, dfs(1, 0) - 1, 1000000007); + } + +private: + long long pow(long long a, int n, int mod) { + long long res = 1; + while (n > 0) { + if (n & 1) { + res = res * a % mod; + } + a = a * a % mod; + n >>= 1; + } + return res; + } +}; ``` #### Go ```go +func assignEdgeWeights(edges [][]int) int { + const mod = 1_000_000_007 + + n := len(edges) + 1 + g := make([][]int, n+1) + + for _, e := range edges { + u, v := e[0], e[1] + g[u] = append(g[u], v) + g[v] = append(g[v], u) + } + + var dfs func(int, int) int + dfs = func(i, fa int) int { + res := 0 + for _, j := range g[i] { + if j != fa { + res = max(res, dfs(j, i)+1) + } + } + return res + } + + return pow(2, dfs(1, 0)-1, mod) +} + +func pow(a, n, mod int) int { + res := 1 + for n > 0 { + if n&1 > 0 { + res = res * a % mod + } + a = a * a % mod + n >>= 1 + } + return res +} +``` +#### TypeScript + +```ts +function assignEdgeWeights(edges: number[][]): number { + const mod = 1_000_000_007; + const n = edges.length + 1; + const g: number[][] = Array.from({ length: n + 1 }, () => []); + + for (const [u, v] of edges) { + g[u].push(v); + g[v].push(u); + } + + const dfs = (i: number, fa: number): number => { + let res = 0; + for (const j of g[i]) { + if (j !== fa) { + res = Math.max(res, dfs(j, i) + 1); + } + } + return res; + }; + + const pow = (a: number, n: number, mod: number): number => { + let res = 1n; + let x = BigInt(a); + const m = BigInt(mod); + + while (n > 0) { + if (n & 1) { + res = (res * x) % m; + } + x = (x * x) % m; + n >>= 1; + } + + return Number(res); + }; + + return pow(2, dfs(1, 0) - 1, mod); +} ``` diff --git a/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.cpp b/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.cpp new file mode 100644 index 0000000000000..81c16616ed48a --- /dev/null +++ b/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.cpp @@ -0,0 +1,39 @@ +class Solution { +public: + int assignEdgeWeights(vector>& edges) { + int n = edges.size() + 1; + vector> g(n + 1); + + for (auto& e : edges) { + int u = e[0]; + int v = e[1]; + g[u].push_back(v); + g[v].push_back(u); + } + + auto dfs = [&](this auto&& dfs, int i, int fa) -> int { + int res = 0; + for (int j : g[i]) { + if (j != fa) { + res = max(res, dfs(j, i) + 1); + } + } + return res; + }; + + return pow(2, dfs(1, 0) - 1, 1000000007); + } + +private: + long long pow(long long a, int n, int mod) { + long long res = 1; + while (n > 0) { + if (n & 1) { + res = res * a % mod; + } + a = a * a % mod; + n >>= 1; + } + return res; + } +}; \ No newline at end of file diff --git a/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.go b/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.go new file mode 100644 index 0000000000000..50943f632a65b --- /dev/null +++ b/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.go @@ -0,0 +1,37 @@ +func assignEdgeWeights(edges [][]int) int { + const mod = 1_000_000_007 + + n := len(edges) + 1 + g := make([][]int, n+1) + + for _, e := range edges { + u, v := e[0], e[1] + g[u] = append(g[u], v) + g[v] = append(g[v], u) + } + + var dfs func(int, int) int + dfs = func(i, fa int) int { + res := 0 + for _, j := range g[i] { + if j != fa { + res = max(res, dfs(j, i)+1) + } + } + return res + } + + return pow(2, dfs(1, 0)-1, mod) +} + +func pow(a, n, mod int) int { + res := 1 + for n > 0 { + if n&1 > 0 { + res = res * a % mod + } + a = a * a % mod + n >>= 1 + } + return res +} \ No newline at end of file diff --git a/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.java b/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.java new file mode 100644 index 0000000000000..04530bafb35e5 --- /dev/null +++ b/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.java @@ -0,0 +1,40 @@ +class Solution { + private List[] g; + + public int assignEdgeWeights(int[][] edges) { + int n = edges.length + 1; + g = new List[n + 1]; + Arrays.setAll(g, k -> new ArrayList<>()); + + for (var e : edges) { + int u = e[0]; + int v = e[1]; + g[u].add(v); + g[v].add(u); + } + + return (int) pow(2, dfs(1, 0) - 1, 1_000_000_007); + } + + private int dfs(int i, int fa) { + int res = 0; + for (int j : g[i]) { + if (j != fa) { + res = Math.max(res, dfs(j, i) + 1); + } + } + return res; + } + + private long pow(long a, int n, int mod) { + long res = 1; + while (n > 0) { + if ((n & 1) != 0) { + res = res * a % mod; + } + a = a * a % mod; + n >>= 1; + } + return res; + } +} \ No newline at end of file diff --git a/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.py b/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.py new file mode 100644 index 0000000000000..5656242aabae5 --- /dev/null +++ b/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.py @@ -0,0 +1,16 @@ +class Solution: + def assignEdgeWeights(self, edges: List[List[int]]) -> int: + def dfs(i: int, fa: int = 0) -> int: + res = 0 + for j in g[i]: + if j != fa: + res = max(res, dfs(j, i) + 1) + return res + + n = len(edges) + 1 + g = [[] for _ in range(n + 1)] + for u, v in edges: + g[u].append(v) + g[v].append(u) + d = dfs(1) + return pow(2, d - 1, 10**9 + 7) diff --git a/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.ts b/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.ts new file mode 100644 index 0000000000000..249ba79b8cb2d --- /dev/null +++ b/solution/3500-3599/3558.Number of Ways to Assign Edge Weights I/Solution.ts @@ -0,0 +1,38 @@ +function assignEdgeWeights(edges: number[][]): number { + const mod = 1_000_000_007; + const n = edges.length + 1; + const g: number[][] = Array.from({ length: n + 1 }, () => []); + + for (const [u, v] of edges) { + g[u].push(v); + g[v].push(u); + } + + const dfs = (i: number, fa: number): number => { + let res = 0; + for (const j of g[i]) { + if (j !== fa) { + res = Math.max(res, dfs(j, i) + 1); + } + } + return res; + }; + + const pow = (a: number, n: number, mod: number): number => { + let res = 1n; + let x = BigInt(a); + const m = BigInt(mod); + + while (n > 0) { + if (n & 1) { + res = (res * x) % m; + } + x = (x * x) % m; + n >>= 1; + } + + return Number(res); + }; + + return pow(2, dfs(1, 0) - 1, mod); +}