🔗 🟡 3067. Count Pairs of Connectable Servers in a Weighted Tree Network 1909

tags: `Biweekly Contest 125` `DFS` `數學(Math)` `組合數學(Combinatorics)` `乘法原理`

題意

給定一個無根的加權樹，樹上總共有 $n$ 個節點，分別表示 $n$ 個伺服器，伺服器從 $0$ 到 $n - 1$ 編號。同時給定一個陣列 $edges$ ，其中 $edges[i] = [a_i, b_i, weight_i]$ 表示節點 $a_i$ 和 $b_i$ 之間有一條雙向邊，邊的權值為 $weight_i$ 。再給定一個整數 $signalSpeed$ 。

如果兩個伺服器 $a$ 、 $b$ 可以透過伺服器 $c$ 連接，則稱伺服器 $a$ 和 $b$ 是通過伺服器 $c$ 可連接的：

$a < b$ ， $a \neq c$ 且 $b \neq c$ 。
從 $c$ 到 $a$ 的距離、以及從 $c$ 到 $b$ 的距離，都可以被 $signalSpeed$ 整除。
從 $c$ 到 $b$ 和 $c$ 到 $a$ 的路徑不共享任何邊。

請你返回一個長度為 $n$ 的整數陣列 $ans$ ，其中 $ans[i]$ 表示通過伺服器 $i$ 可連接 的伺服器對的數目。

思路：DFS + 乘法原理

首先注意到條件 $1.$ 中 $a < b$ ，看似需要考慮 $a$ 和 $b$ 的大小，但實際上不需要。由於伺服器的編號皆不相同，因此對於任何點對 $(u, v)$ ，都只會計算一次。

而條件 $3.$ 中由於不能共享任何邊，因此對於 $u$ 的子節點 $v_0, v_1, \cdots, v_k$ ，我們只需考慮以 $v_i$ 為根的子樹中滿足條件 $2.$ 的節點數。由左至右考慮，確保每個點對只會被計算一次，並使用乘法原理計算。其中 $v_0$ 無法與其他子樹中構成符合條件的點對， $v_1$ 可以與 $v_0$ 構成 $v_0 \times v_1$ 個點對， $v_2$ 可以與 $v_0, v_1$ 構成 $(v_0 + v_1) \times v_2$ 個點對，以此類推。可以令 $pre$ 為前面子樹中符合條件的節點數， $cur$ 為當前子樹中符合條件的節點數，則當前節點 $v_i$ 對 $ans[u]$ 的貢獻為 $pre \times cur$ 。

至於要計算以 $v$ 為根的子樹中符合條件的節點數，可以使用 DFS 計算，令 $ps$ 為當前節點與根節點的距離，則當 $ps$ 能被 $signalSpeed$ 整除時，該節點符合條件，遞迴計算子樹中符合條件的節點數。

複雜度分析

時間複雜度 $\mathcal{O}(n^2)$ ，枚舉所有節點做為根節點，而在每次枚舉中，每個節點都會被訪問一次。
空間複雜度 $\mathcal{O}(n)$ ，存儲樹的 鄰接表(Adjacency list) 。

class Solution:
    def countPairsOfConnectableServers(self, edges: List[List[int]], signalSpeed: int) -> List[int]:
        n = len(edges) + 1 # 注意樹中點的數量是邊數加一
        g = [[] for _ in range(n)]
        for u, v, w in edges:
            g[u].append((v, w))
            g[v].append((u, w))

        def dfs(u: int, fa: int, ps: int) -> int: # ps 是與根節點的距離
            res = 1 if ps % signalSpeed == 0 else 0 # 與根節點距離是 signalSpeed 的倍數的節點數
            for v, w in g[u]:
                if v == fa:
                    continue
                res += dfs(v, u, ps + w)
            return res
        ans = [0] * n

        for u, gu in enumerate(g): # 枚舉每個節點作為根節點(中間點)
            pre = 0 # 前面子樹中符合條件的節點數
            for v, w in gu: 
                cur = dfs(v, u, w) # 以 v 為根的子樹，符合條件的節點數
                ans[u] += pre * cur # 乘法原理
                pre += cur
        return ans

class Solution {
public:
    vector<int> countPairsOfConnectableServers(vector<vector<int>>& edges, int signalSpeed) {
        int n = edges.size() + 1; // 注意樹中點的數量是邊數加一
        vector<vector<pair<int, int>>> g(n);
        for (auto& e : edges) {
            g[e[0]].emplace_back(e[1], e[2]);
            g[e[1]].emplace_back(e[0], e[2]);
        }

        function <int(int, int, int)> dfs = [&](int u, int fa, int ps) -> int { // ps 是與根節點的距離
            int res = ps % signalSpeed == 0; // 與根節點距離是 signalSpeed 的倍數的節點數
            for (auto& vw: g[u]) {
                int v = vw.first, w = vw.second;
                if (v == fa) continue;
                res += dfs(v, u, ps + w);
            }
            return res;
        };

        vector<int> ans(n);
        for (int u = 0; u < n; u++) { // 枚舉每個節點作為根節點
            int pre = 0; // 前面子樹中符合條件的節點數
            for (auto& vw: g[u]) {
                int v = vw.first, w = vw.second;
                int cur = dfs(v, u, w); // 以 v 為根的子樹，符合條件的節點數
                ans[u] += pre * cur; // 乘法原理
                pre += cur;
            }
        }
        return ans;
    }
};