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

Array
Tree
DFS

TIME

O(n²)

SPACE

O(n)

Intuition

For each server c, count valid servers in each branch (subtree), then count pairs across different branches. Two servers are connectable if they're in different branches and both have distances divisible by signalSpeed.

Solution

Build adjacency list: Create graph from edges with weights.
For each potential middle server c:
- For each neighbor (branch), DFS to count servers with distance divisible by signalSpeed
- Store count for each branch
Count pairs across branches:
- Use prefix sum approach: for each branch count, multiply by sum of all previous branch counts
- This counts all pairs (a, b) where a and b are in different branches
DFS helper:
- Traverse subtree rooted at node (excluding parent)
- Count nodes where distance % signalSpeed == 0

Why Prefix Sum for Pairs?

If branches have counts [c1, c2, c3]:

Pairs = c1c2 + c1c3 + c2*c3
Using prefix: (0)*c1 + (c1)*c2 + (c1+c2)*c3

Complexity

Time: O(n²), for each node, DFS entire tree
Space: O(n), adjacency list and recursion stack

Edge Cases

Linear tree (path graph) → left count × right count
Star graph → all pairs through center
signalSpeed = 1 → all distances valid
No valid pairs → all zeros

Code

from typing import List
from collections import defaultdict


class Solution:
    def countPairsOfConnectableServers(self, edges: List[List[int]], signalSpeed: int) -> List[int]:
        """
        3067. Count Pairs of Connectable Servers in a Weighted Tree Network
        Time: O(n²)
        Space: O(n)
        """
        n = len(edges) + 1
        graph = defaultdict(list)

        for u, v, w in edges:
            graph[u].append((v, w))
            graph[v].append((u, w))

        result = [0] * n

        def dfs(node: int, parent: int, dist: int) -> int:
            count = 1 if dist % signalSpeed == 0 else 0
            for nei, w in graph[node]:
                if nei == parent:
                    continue
                count += dfs(nei, node, dist + w)
            return count

        for c in range(n):
            branch_counts = []

            for nei, w in graph[c]:
                branch_counts.append(dfs(nei, c, w))

            total = 0
            prefix = 0
            for cnt in branch_counts:
                total += prefix * cnt
                prefix += cnt

            result[c] = total

        return result

import java.util.ArrayList;
import java.util.List;

/**
 * 3067. Count Pairs of Connectable Servers in a Weighted Tree Network
 * Time: O(n²)
 * Space: O(n)
 */
class Solution {
    private List<int[]>[] graph;
    private int signalSpeed;

    public int[] countPairsOfConnectableServers(int[][] edges, int signalSpeed) {
        int n = edges.length + 1;
        this.signalSpeed = signalSpeed;

        graph = new ArrayList[n];
        for (int i = 0; i < n; i++) {
            graph[i] = new ArrayList<>();
        }

        for (int[] e : edges) {
            graph[e[0]].add(new int[]{e[1], e[2]});
            graph[e[1]].add(new int[]{e[0], e[2]});
        }

        int[] result = new int[n];

        for (int c = 0; c < n; c++) {
            List<Integer> branchCounts = new ArrayList<>();

            for (int[] nei : graph[c]) {
                branchCounts.add(dfs(nei[0], c, nei[1]));
            }

            int total = 0;
            int prefix = 0;
            for (int cnt : branchCounts) {
                total += prefix * cnt;
                prefix += cnt;
            }

            result[c] = total;
        }

        return result;
    }

    private int dfs(int node, int parent, int dist) {
        int count = (dist % signalSpeed == 0) ? 1 : 0;

        for (int[] nei : graph[node]) {
            if (nei[0] == parent) continue;
            count += dfs(nei[0], node, dist + nei[1]);
        }

        return count;
    }
}