LeetCode 解題筆記：3558. Number of Ways to Assign Edge Weights I

作者：王一哲
日期：2026年6月11日


LeetCode 題目連結：3558. Number of Ways to Assign Edge Weights I

解題想法

中等難度題目。我先用接鄰矩陣 $adj$ 存樹，題目規定根節點編號為 $1$，父節點編號小於子節點編號，但是測資 $edges$ 之中的節點編號可能會是 $u > v$，如果遇到這個狀況要先交換 $u, v$ 之後再存入 $adj$之中。接下來寫一個自訂函式 $dfs$，從根節點 $1$ 開始往下走找最大深度 max_depth。最後答案為 2**(max_depth - 1) % 1000000007。由於答案很大而且需要取餘數，如果用 Python 解題可以直接用內建的 pow，但如果用 C++ 解題要自己寫快速冪。 證明答案的過程如下。假設有 $m$ 條邊，且 $m \geq 1$，選擇其中奇數條邊設定為 $1$ 可以符合題目的要求，因此組合數為 $$ \binom{m}{1} + \binom{m}{3} + \binom{m}{5} + \dots \binom{m}{m} = 2^{m-1} $$ 可以透過二項式定理 (binomial theorem) 導出這個結果。根據二項式展開式 $$ (1 + x)^m = \binom{m}{0} + \binom{m}{1}x + \binom{m}{2}x^2 + \dots + \binom{m}{m}x^m $$ 為了求出奇數項的總和，先對上式代入 $x = 1$，計算從 $m$ 條邊中隨意選擇任意數量的所有可能組合數 $$ 2^m = \binom{m}{0} + \binom{m}{1} + \binom{m}{2} + \binom{m}{3} + \dots + \binom{m}{m} $$ 再代入 $x = -1$，選擇偶數條邊的項，選擇奇數條邊的項變成負號。 $$ 0 = \binom{m}{0} - \binom{m}{1} + \binom{m}{2} - \binom{m}{3} + \dots + (-1)^m \binom{m}{m} $$ 第 1 式減去第 2 式，偶數項會被完全抵消，而奇數項會變成兩倍： $$ 2^m - 0 = 2 \times \left[ \binom{m}{1} + \binom{m}{3} + \binom{m}{5} + \dots \right] $$ 最後，將等式兩邊同除以 2，即可得到選擇奇數條邊的組合數總和： $$ \binom{m}{1} + \binom{m}{3} + \binom{m}{5} + \dots = \frac{2^m}{2} = 2^{m-1} $$

Python 程式碼

Runtime: 267 ms, beats 84.95%. Memory: 104.95 MB, beats 38.71%.

class Solution:
    def assignEdgeWeights(self, edges: List[List[int]]) -> int:
        n = len(edges) + 1  # n 個節點，n-1 條邊
        adj = [[] for _ in range(n+1)]  # 用接鄰矩陣存樹
        for u, v in edges:
            if u > v:
                u, v = v, u
            adj[u].append(v)

        # 用 dfs 找最大深度
        max_depth = 0

        def dfs(u, depth):
            nonlocal max_depth
            # 遞迴出口
            if not adj[u]:
                max_depth = max(max_depth, depth)
            # 繼續往下走
            for v in adj[u]:
                dfs(v, depth + 1)
        
        # 呼叫 dfs
        dfs(1, 0)  # 從根節點 1、深度 0 開始
        
        # 答案是 2**(max_depth - 1) % (10**9 + 7)
        return pow(2, max_depth - 1, 10**9 + 7)

C++ 程式碼

Runtime: 287 ms, beats 61.07%. Memory: 332.90 MB, beats 68.46%.

class Solution {
public:
    /* 快速冪，計算次方時取餘數 */
    long long mypow(long long a, long long b, const int MOD) {
        long long r = 1;
        while(b) {
            if (b&1) {
                r = (r * a) % MOD;
            }
            a = (a * a) % MOD;
            b >>= 1; 
        }
        return r;
    }

    vector<vector<int>> adj;  // 用接鄰矩陣存樹
    int max_depth = 0;  // 最大深度

    void dfs(int u, int depth) {
        // 遞迴出口
        if (adj[u].empty()) {
            max_depth = max(max_depth, depth);
        }
        // 繼續往下走
        for(int v : adj[u]) {
            dfs(v, depth + 1);
        }
    }

    int assignEdgeWeights(vector<vector<int>>& edges) {
        int n = (int)edges.size() + 1;  // n 個節點，n-1 條邊
        adj.assign(n+1, vector<int> (0));  // 調整 adj 的長度
        for(auto it : edges) {
            int u = it[0], v = it[1];
            if (u > v) swap(u, v);  // 如果 u 大於 v，交換
            adj[u].push_back(v);
        }

        // 呼叫 dfs 找最大深度
        dfs(1, 0);  // 從根節點 1、深度 0 開始
        
        // 答案是 2**(max_depth - 1) % (10**9 + 7)
        return mypow(2, max_depth - 1, 1000000007);
    }
};