3067. 可連結なサーバーペアの数を重み付き木ネットワークで統計する

3067. 在带权树网络中统计可连接服务器对数目 #

与えられた根付き重み付き木があります。木には合計で $n$ 個のノードがあり、それぞれ $n$ 個のサーバーを表します。サーバーは $0$ から $n - 1$ まで番号が付けられています。また、配列 $edges$ が与えられます。ここで、 $edges[i] = [a_i, b_i, weight_i]$ はノード $a_i$ と $b_i$ の間に双方向のエッジがあり、エッジの重みは $weight_i$ です。さらに、整数 $signalSpeed$ も与えられます。

2 つのサーバー $a$ 、 $b$ 、 $c$ が以下の条件を満たす場合、サーバー $a$ と $b$ はサーバー $c$ を介して接続可能であると言います。

$a < b$ 、 $a\ne c$ かつ $b\ne c$ 。
$c$ から $a$ への距離は $signalSpeed$ で割り切れる。
$c$ から $b$ への距離は $signalSpeed$ で割り切れる。
$c$ から $b$ へのパスと $c$ から $a$ へのパスには共通のエッジが存在しない。

長さ $n$ の整数配列 $count$ を返し、 $count[i]$ がサーバー $i$ を介して接続可能なサーバーペアの数を表すようにしてください。

例 1：

入力：edges = [[0,1,1],[1,2,5],[2,3,13],[3,4,9],[4,5,2]], signalSpeed = 1
出力：[0,4,6,6,4,0]
説明：signalSpeed が 1 のため、count[c] は c を始点とし、共通のエッジがないパスの数です。
入力グラフでは、count[c] はサーバー c の左側のサーバーの数と右側のサーバーの数の積です。

例 2：

入力：edges = [[0,6,3],[6,5,3],[0,3,1],[3,2,7],[3,1,6],[3,4,2]], signalSpeed = 3
出力：[2,0,0,0,0,0,2]
説明：サーバー 0 を介して、接続可能なサーバーペア (4, 5) と (4, 6) の2つがあります。
サーバー 6 を介して、接続可能なサーバーペア (4, 5) と (0, 5) の2つがあります。
他のサーバーペアはすべてサーバー 0 または 6 を介して接続する必要があるため、それ以外のサーバーペアの数はすべて 0 です。

ヒント：

$2 <= n <= 1000$
$edges.length == n - 1$
$edges[i].length == 3$
$0 <= a_i, b_i < n$
$edges[i] = [a_i, b_i, weight_i]$
$1 <= weight_i <= 10^6$
$1 <= signalSpeed <= 10^6$
入力は $edges$ が有効な木を構成していることを保証します。

根の列挙#

$a < b$ 、 $a \ne c$ かつ $b \ne c$ 。
$c$ から $a$ への距離は $signalSpeed$ で割り切れる。
$c$ から $b$ への距離は $signalSpeed$ で割り切れる。
$c$ から $b$ へのパスと $c$ から $a$ へのパスには共通のエッジが存在しない。

接続可能な条件から、次の情報を得ることができます。

a と b の位置は交換可能であり、差が存在することを満たすだけです。（重複を防ぐためとも理解できます）

3067. 在带权树网络中统计可连接服务器对数目_1

c はこの無向非巡回グラフを 2 つの部分に分けるノードです。c から a へのパスと c から b へのパスには共通のエッジがないためです。

3067. 在带权树网络中统计可连接服务器对数目_2

以下のケースを分析します：

サブツリー 1 を探索：

サブツリー 2 を探索：

サブツリー 3 を探索：

各探索で条件を満たす結果を合計し、現在の探索の根ノードの接続可能なノードペアとします。

コードの実装：

class Solution {
public:
    vector<int> countPairsOfConnectableServers(vector<vector<int>> &edges, int signalSpeed) {
        int n = edges.size() + 1; // 辺の数はノードの数より1少ない
        // {root, [<child1, cost1>, <child2, cost2>, ..., <childn, costn>]}
        vector<vector<pair<int, int>>> graph(n);

        for (const auto &edge : edges) {
            graph[edge[0]].push_back({edge[1], edge[2]});
            graph[edge[1]].push_back({edge[1], edge[2]});
        }

        /// @brief 
        /// @param root 根ノード、このパラメータを渡すことで、サブツリーの根ノードがこのノードに向かうのを防ぐ
        /// @param child 子ノード
        /// @param remainder root から child までの距離
        /// @return root ノード、child を1つのサブツリーとして、signalSpeed で割り切れる到達パスを持つノードの数を返す
        function<int(int, int, int)> dfs = [&](int root, int child, int remainder) -> int {
            int res = 0;
            if (remainder == 0) {
                ++res;
            }

            // 子ノードのすべてのサブツリーを探索
            for (const auto &[j, cost] : graph[child]) {
                // ルートノードのパスを除外
                if (j != root) {
                    res += dfs(child, j, (remainder + cost) % signalSpeed);
                }
            }
            return res;
        };

        vector<int> res(n);

        // ルートノードを探索
        for (int i = 0; i < n; ++i) {
            int pre = 0; // ルートノードのサブツリーの接続可能なペア数

            // ルートノードのすべてのサブツリーを探索
            for (auto &[j, cost] : graph[i]) {
                // 現在のサブツリーが満たす条件のノード数
                int cnt = dfs(i, j, cost % signalSpeed);
                res[i] += pre * cnt;
                pre += cnt;
            }
        }
        return res;
    }
};