Count Pairs of Connectable Servers in a Weighted Tree Network
Source: leetcode
Difficulty: Medium
Tags: array, tree, depth-first-search

You are given an unrooted weighted tree with `n` vertices representing servers numbered from `0` to `n - 1`, an array `edges` where `edges[i] = [ai, bi, weighti]` represents a bidirectional edge between vertices `ai` and `bi` of weight `weighti`. You are also given an integer `signalSpeed`.


Two servers `a` and `b` are **connectable** through a server `c` if:


- `a < b`, `a != c` and `b != c`.
- The distance from `c` to `a` is divisible by `signalSpeed`.
- The distance from `c` to `b` is divisible by `signalSpeed`.
- The path from `c` to `b` and the path from `c` to `a` do not share any edges.


Return *an integer array* `count` *of length* `n` *where* `count[i]` *is the **number** of server pairs that are **connectable** through* *the server* `i`.


 


**Example 1:**


![](https://assets.leetcode.com/uploads/2024/01/21/example22.png)

```

Input: edges = [[0,1,1],[1,2,5],[2,3,13],[3,4,9],[4,5,2]], signalSpeed = 1
Output: [0,4,6,6,4,0]
Explanation: Since signalSpeed is 1, count[c] is equal to the number of pairs of paths that start at c and do not share any edges.
In the case of the given path graph, count[c] is equal to the number of servers to the left of c multiplied by the servers to the right of c.

```

**Example 2:**


![](https://assets.leetcode.com/uploads/2024/01/21/example11.png)

```

Input: edges = [[0,6,3],[6,5,3],[0,3,1],[3,2,7],[3,1,6],[3,4,2]], signalSpeed = 3
Output: [2,0,0,0,0,0,2]
Explanation: Through server 0, there are 2 pairs of connectable servers: (4, 5) and (4, 6).
Through server 6, there are 2 pairs of connectable servers: (4, 5) and (0, 5).
It can be shown that no two servers are connectable through servers other than 0 and 6.

```

 


**Constraints:**


- `2 <= n <= 1000`
- `edges.length == n - 1`
- `edges[i].length == 3`
- `0 <= ai, bi < n`
- `edges[i] = [ai, bi, weighti]`
- `1 <= weighti <= 106`
- `1 <= signalSpeed <= 106`
- The input is generated such that `edges` represents a valid tree.




C++ solution:
class Solution {
public:
    vector<int> countPairsOfConnectableServers(vector<vector<int>>& edges, int signalSpeed) {
        ios::sync_with_stdio(false);cin.tie(NULL);cout.tie(NULL);
        int n=edges.size()+1;
        typedef pair<int,int> ii;
        vector<vector<ii>> g(n);
        for(auto& e:edges){
            g[e[0]].emplace_back(e[1], e[2]);
            g[e[1]].emplace_back(e[0], e[2]);
        }

        auto dfs=[&](int u, int p, int d, auto&& dfs)->int{
            int res=0;
            if(d%signalSpeed==0){
                res++;
            }
            for(auto& [v,w]: g[u]){
                if(v!=p){
                    res+=dfs(v,u,d+w, dfs);
                }
            }
            return res;
        };

        vector<int> ans(n);
        for(int i=0; i<n; i++){
            int m=g[i].size();
            if(m==1){
                ans[i]=0;
                continue;
            }
            int sum=0;
            ans[i]=0;
            for(auto& [v,w]:g[i]){
                int r=dfs(v,i,w, dfs);
                ans[i]+=sum*r;
                sum+=r;
            }
        }
        return ans;
    }
};

Go solution:
var g [][][]int
var sp int

func countPairsOfConnectableServers(edges [][]int, signalSpeed int) []int {
    n:=len(edges)+1
    sp=signalSpeed
    g=make([][][]int, n)
    for i:=range g{
        g[i]=[][]int{}
    }
    for _,e:=range edges{
        g[e[0]]=append(g[e[0]], []int{e[1],e[2]})
        g[e[1]]=append(g[e[1]], []int{e[0],e[2]})
    }
    ans:=make([]int, 0, n)
    for i:=range g{
        sqSum,sum:=0,0
        for _,e:=range g[i]{
            count:=dfs(i,e[0],e[1])
            sqSum+=count*count
            sum+=count
        }
        ans=append(ans, (sum*sum-sqSum)/2)
    }
    return ans
}

func dfs(p, i, d int) int {
    ans:=0
    if d%sp == 0{
        ans++
    }
    for _,e:=range g[i]{
        if e[0]==p {continue}
        ans+=dfs(i,e[0],d+e[1])
    }
    return ans
}

Java solution:
class Solution {
    ArrayList<ArrayList<int[]>> graph;
    int mod;
    public int[] countPairsOfConnectableServers(int[][] edges, int signalSpeed) {
        int n = edges.length+1; 
        int res[] = new int[n];
        mod = signalSpeed;
        
        graph = new ArrayList<ArrayList<int[]>>();
        
        for(int i = 0; i<n; i++)graph.add(new ArrayList<int[]>());
        
        for(int e[] : edges){
            graph.get(e[0]).add(new int[]{e[1],e[2]});
            graph.get(e[1]).add(new int[]{e[0],e[2]});
        }
        
        for(int i = 0; i<n; i++){
            res[i] = fun(i);
        }
        
        return res;
    }
    
    
    int fun(int i){
        int res = 0,cur = 0; 
        
        for(int j[] : graph.get(i)){
            int t = dfs(j[0],i,j[1]);
            res += t*cur;
            cur += t;
        }
        
        return res;
    }
    
    int dfs(int node, int par, int wt){
        int res = wt%mod == 0 ? 1: 0;
        
        for(int i[] : graph.get(node)){
            if(i[0] == par)continue;
            res += dfs(i[0],node,wt+i[1]);
        }
        
        return res;
    }
}

JavaScript solution:
/**
 * @param {number[][]} edges
 * @param {number} signalSpeed
 * @return {number[]}
 */
function sumOfPairwiseProducts(arr) {
  let sum = 0;
  const n = arr.length;

  for (let i = 0; i < n; i++) {
    for (let j = i + 1; j < n; j++) {
      sum += arr[i] * arr[j];
    }
  }

  return sum;
}
var countPairsOfConnectableServers = function (edges, signalSpeed) {
  let servers = new Set();
  for (let [a, b, w] of edges) {
    servers.add(a);
    servers.add(b);
  }
  let n = servers.size;
  let adjList = Array(n)
    .fill(0)
    .map(() => []);
  let visited = new Set();
  for (let [a, b, w] of edges) {
    adjList[a].push([b, w]);
    adjList[b].push([a, w]);
  }
  let dfs = (node, distance) => {
    if (visited.has(node)) {
      return 0;
    }
    visited.add(node);
    let countOfNodesWithDistanceDivBySpeed = 0;
    let neighbours = adjList[node];
    for (let [val, w] of neighbours) {
      countOfNodesWithDistanceDivBySpeed += dfs(val, distance + w);
    }
    if (distance % signalSpeed === 0) {
      countOfNodesWithDistanceDivBySpeed++;
    }
    return countOfNodesWithDistanceDivBySpeed;
  };
  let res = [];
  for (let c = 0; c < n; c++) {
    let neighbours = adjList[c];
    let neis = neighbours
      .filter((n) => n[1] % signalSpeed === 0)
      .map((e) => e[0]);
    if (neighbours.length <= 1) {
      res.push(0);
      continue;
    }
    let divisibles = [];
    for (let [val, w] of neighbours) {
      visited = new Set();
      visited.add(c);
      let countOfNodesWithDistanceDivBySpeed = dfs(val, w);
      divisibles.push(countOfNodesWithDistanceDivBySpeed);
    }
    res.push(sumOfPairwiseProducts(divisibles));
  }
  return res;
};

Python solution:
class Solution:
    def countPairsOfConnectableServers(self, edges: List[List[int]], k: int) -> List[int]:
        l=len(edges)+1;g=[[] for _ in range(l)]
        for a, b, w in edges:
            g[a].append((b, w))
            g[b].append((a, w))
        
        ans=[0]*l

        def rec(v, par):
            lst=[];lst2=[];s1=0;s2=0;dv=defaultdict(int);d2v=defaultdict(list)
            for v2, w in g[v]:
                if v2!=par:
                    dv2, d2v2=rec(v2, v)
                    lst.append(dv2)
                    lst2.append(d2v2)
                    for val, ct in dv2.items():
                        nval=(val+w)%k
                        if nval==0:
                            s1+=ct
                            s2+=ct*ct
                        dv[nval]+=ct                        
                    for val2, ls2 in d2v2.items():
                        nval2=(val2+w)%k
                        d2v[nval2].extend(ls2)
            
            dv[0]+=1
            ans[v]+=(s1*s1-s2)>>1

            i=0
            for v2, w in g[v]:
                if v2!=par:
                    dv2=lst[i]
                    d2v2=lst2[i]
                    i+=1
                    for val2, ls2 in d2v2.items():
                        nval2=(-val2-w)%k
                        if nval2 in dv:
                            ct=dv[nval2]-dv2[(-val2-(w<<1))%k]
                            for v3, ct3 in ls2:
                                ans[v3]+=ct3*ct

            if s1:
                d2v[0].append((v, s1))
            
            return dv, d2v

        rec(l>>1, -1)
        return ans

Ruby solution:
def count_pairs_of_connectable_servers(edges, signal_speed)
  n = edges.size + 1
  adj = Array.new(n) { [] }
  edges.each do |a, b, w|
    adj[a] << [b, w]
    adj[b] << [a, w]
  end

  ans = Array.new(n, 0)

  (0...n).each do |c|
    counts = []
    adj[c].each do |v, w|
      count = dfs(v, c, w, signal_speed, adj)
      counts << count
    end

    sum = counts.sum
    sum_sq = counts.sum { |x| x * x }
    ans[c] = (sum * sum - sum_sq) / 2
  end

  ans
end

def dfs(current, parent, current_dist, signal_speed, adj)
  count = 0
  count += 1 if current_dist % signal_speed == 0

  adj[current].each do |neighbor, weight|
    if neighbor != parent
      count += dfs(neighbor, current, current_dist + weight, signal_speed, adj)
    end
  end

  count
end