Count Pairs of Connectable Servers in a Weighted Tree Network

medium tree dfs

Problem

You have an unrooted weighted tree with n servers. Two servers a and b are connectable through a third server c (with a ≠ c, b ≠ c) if the distance from a to c and the distance from b to c are both divisible by signalSpeed, and a and b lie on different branches of c. For every server c, return the number of such connectable pairs.

Inputedges = [[0,1,1],[1,2,5],[2,3,13],[3,4,9],[4,5,2]], signalSpeed = 1

Output[0,4,6,6,4,0]

With signalSpeed 1 every distance is divisible, so each interior server pairs up the branch sizes around it; the endpoints have a single branch and contribute 0.

edges as [u,v,weight] list

signalSpeed

def count_pairs(edges, signalSpeed):
    n = len(edges) + 1
    g = [[] for _ in range(n)]
    for a, b, w in edges:
        g[a].append((b, w)); g[b].append((a, w))
    def dfs(node, parent, dist):
        cnt = 1 if dist % signalSpeed == 0 else 0
        for nx, w in g[node]:
            if nx != parent:
                cnt += dfs(nx, node, dist + w)
        return cnt
    res = [0] * n
    for c in range(n):
        counts = [dfs(nx, c, w) for nx, w in g[c]]
        pairs, prefix = 0, 0
        for cc in counts:
            pairs += prefix * cc
            prefix += cc
        res[c] = pairs
    return res

function countPairs(edges, signalSpeed) {
  const n = edges.length + 1;
  const g = Array.from({length: n}, () => []);
  for (const [a, b, w] of edges) { g[a].push([b, w]); g[b].push([a, w]); }
  function dfs(node, parent, dist) {
    let cnt = dist % signalSpeed === 0 ? 1 : 0;
    for (const [nx, w] of g[node])
      if (nx !== parent) cnt += dfs(nx, node, dist + w);
    return cnt;
  }
  const res = new Array(n).fill(0);
  for (let c = 0; c < n; c++) {
    let pairs = 0, prefix = 0;
    for (const [nx, w] of g[c]) {
      const cc = dfs(nx, c, w);
      pairs += prefix * cc;
      prefix += cc;
    }
    res[c] = pairs;
  }
  return res;
}

import java.util.*;
class Solution {
    int speed;
    List<int[]>[] g;
    public int[] countPairsOfConnectableServers(int[][] edges, int signalSpeed) {
        int n = edges.length + 1; speed = signalSpeed;
        g = new List[n];
        for (int i = 0; i < n; i++) g[i] = new ArrayList<>();
        for (int[] e : edges) { g[e[0]].add(new int[]{e[1], e[2]}); g[e[1]].add(new int[]{e[0], e[2]}); }
        int[] res = new int[n];
        for (int c = 0; c < n; c++) {
            long pairs = 0, prefix = 0;
            for (int[] nb : g[c]) {
                long cc = dfs(nb[0], c, nb[1]);
                pairs += prefix * cc; prefix += cc;
            }
            res[c] = (int) pairs;
        }
        return res;
    }
    long dfs(int node, int parent, int dist) {
        long cnt = dist % speed == 0 ? 1 : 0;
        for (int[] nb : g[node]) if (nb[0] != parent) cnt += dfs(nb[0], node, dist + nb[1]);
        return cnt;
    }
}

class Solution {
public:
    int speed;
    vector<vector<pair<int,int>>> g;
    long long dfs(int node, int parent, int dist) {
        long long cnt = dist % speed == 0 ? 1 : 0;
        for (auto& nb : g[node]) if (nb.first != parent) cnt += dfs(nb.first, node, dist + nb.second);
        return cnt;
    }
    vector<int> countPairsOfConnectableServers(vector<vector<int>>& edges, int signalSpeed) {
        int n = edges.size() + 1; speed = signalSpeed;
        g.assign(n, {});
        for (auto& e : edges) { g[e[0]].push_back({e[1], e[2]}); g[e[1]].push_back({e[0], e[2]}); }
        vector<int> res(n, 0);
        for (int c = 0; c < n; c++) {
            long long pairs = 0, prefix = 0;
            for (auto& nb : g[c]) {
                long long cc = dfs(nb.first, c, nb.second);
                pairs += prefix * cc; prefix += cc;
            }
            res[c] = (int) pairs;
        }
        return res;
    }
};

Explanation

Fix one server c as the meeting point. Every other server reaches c through exactly one of c's direct branches. A pair is connectable through c when both servers sit in different branches and both of their distances to c are multiples of signalSpeed.

So for each branch of c we run a DFS that counts how many servers in that branch have a distance to c divisible by signalSpeed. Call those branch totals cc.

The number of valid pairs is the number of ways to pick two qualifying servers from two different branches. If we process branches left to right keeping a running prefix of qualifying servers seen so far, each new branch contributes prefix * cc new cross-branch pairs.

We repeat this for every server as the center. Leaf servers have a single branch, so there is no second branch to pair with and their answer is 0.

Example: with signalSpeed = 1 every distance is divisible by 1. Server 2 sees branches of sizes 2 and 3, giving 2 * 3 = 6 pairs, matching the output [0,4,6,6,4,0].

Time: O(n²) Space: O(n)