Steward

分享是一種喜悅、更是一種幸福

程式語言 - LeetCode - C++ - 3558. Number of Ways to Assign Edge Weights I

---

題目：


解答：

class Solution {
public:
    int assignEdgeWeights(vector<vector<int>>& edges) {
        int n = edges.size() + 1;
        vector<vector<int>> g(n + 1);

        for (auto e : edges) {
            g[e[0]].push_back(e[1]);
            g[e[1]].push_back(e[0]);
        }

        int mx = 0;

        auto dfs = [&](this auto&& dfs, int u, int parent, int depth) -> void {
            mx = max(mx, depth);

            for (int n : g[u]) {
                if (n == parent) {
                    continue;
                }

                dfs(n, u, depth + 1);
            }
        };

        constexpr int MOD = 1e9 + 7;

        auto mod = [](long long a, long long b) -> int {
            long long ans = 1;

            while (b) {
                if (b & 1) {
                    ans = ans * a % MOD;
                }

                a = a * a % MOD;
                b >>= 1;
            }

            return ans;
        };

        dfs(1, 0, 0);

        return mod(2, mx - 1);
    }
};