Count Pairs of Connectable Servers in a Weighted Tree Network

Med

#2723Time: O(N^3). The main loop runs `N` times. Inside, the traversal takes O(N). The nested loops to check all pairs `(a, b)` take O(N^2). Thus, the total complexity is O(N * (N + N^2)) = O(N^3), which is too slow for the given constraints.Space: O(N) to store the adjacency list, a distance array, and a first-hop array for each central server.2 companies

Algorithms

Depth-First Search

Data structures

Array Tree

Companies

UBS thoughtspot

Prompt

You are given an unrooted weighted tree with n vertices representing servers numbered from 0 to n - 1, an array edges where edges[i] = [a_i, b_i, weight_i] represents a bidirectional edge between vertices a_i and b_i of weight weight_i. 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:

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:

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 <= a_i, b_i < n
edges[i] = [a_i, b_i, weight_i]
1 <= weight_i <= 10⁶
1 <= signalSpeed <= 10⁶
The input is generated such that edges represents a valid tree.

Approaches

2 approaches with complexity analysis and trade-offs.

This is a direct, brute-force approach. For each server in the network, we consider it as the central server c. Then, we examine every possible pair of other servers (a, b) to see if they are connectable through c.

Algorithm

Build an adjacency list from the edges array.
Initialize an answer array result of size n.
For each server c from 0 to n-1:
- Perform a graph traversal (like BFS) starting from c to calculate distances and the first neighbor on the path to every other node. Store these in dist and firstHop arrays.
- Initialize pairCount = 0.
- Iterate a from 0 to n-1.
- Iterate b from a + 1 to n-1.
- If a == c or b == c, skip this pair.
- Check if dist[a] % signalSpeed == 0, dist[b] % signalSpeed == 0, and firstHop[a] != firstHop[b].
- If all true, increment pairCount.
- Set result[c] = pairCount.
Return result.

Walkthrough

To check if a pair (a, b) is connectable through c, we must verify all four conditions given in the problem. The main computational tasks are finding the distances from c to a and b, and ensuring their paths from c are edge-disjoint.

A single traversal like Breadth-First Search (BFS) or Depth-First Search (DFS) starting from c can efficiently compute the distances to all other nodes in O(N) time. During this same traversal, we can also determine the "first hop" for each node, which is the neighbor of c on the unique path from c to that node. The path disjointness condition is met if and only if a and b have different first hops from c.

The overall algorithm is as follows:

Build an adjacency list representation of the tree.
For each server c from 0 to n-1: a. Run a traversal (e.g., BFS) from c to compute distances[i] and first_hops[i] for all other nodes i. b. Initialize a counter for connectable pairs to zero. c. Iterate through all pairs of servers (a, b) with a < b. d. If a or b is the same as c, skip. e. Check if distances[a] and distances[b] are divisible by signalSpeed, and if first_hops[a] is different from first_hops[b]. f. If all conditions hold, increment the pair counter. g. After checking all pairs, store the result for server c.
Return the final counts.

Complexity

Time

O(N^3). The main loop runs `N` times. Inside, the traversal takes O(N). The nested loops to check all pairs `(a, b)` take O(N^2). Thus, the total complexity is O(N * (N + N^2)) = O(N^3), which is too slow for the given constraints.

Space

O(N) to store the adjacency list, a distance array, and a first-hop array for each central server.

Trade-offs

Pros

Conceptually simple and directly follows the problem definition.

Cons

The O(N^3) time complexity makes it infeasible for the given constraints.

This approach improves upon the brute-force method by changing the way we count pairs. Instead of checking every pair (a, b), for each potential central server c, we first count how many valid nodes exist in each of its branches. A "branch" corresponds to a subtree attached to one of c's neighbors. If a branch i has count_i valid nodes, and another branch j has count_j valid nodes, they contribute count_i * count_j pairs to the total.

Algorithm

Build an adjacency list to represent the tree.
Initialize an answer array result of size n with zeros.
Loop for i from 0 to n-1 (treating i as the central server c):
- If i has fewer than two neighbors, it cannot connect any pair, so result[i] remains 0. Continue.
- Initialize totalPairs = 0 and prefixCount = 0.
- For each neighbor v of i:
  - Start a traversal (e.g., DFS) from v, with i as the parent to avoid going back.
  - The initial distance for the traversal is the weight of the edge (i, v).
  - The traversal counts nodes x in the subtree where distance(i, x) is divisible by signalSpeed. Let this be countInSubtree.
  - Add countInSubtree * prefixCount to totalPairs.
  - Update prefixCount by adding countInSubtree.
- Set result[i] = totalPairs.
Return result.

Walkthrough

The algorithm iterates through each server i and considers it as the central server c. For c, the paths to any two nodes a and b are edge-disjoint if and only if a and b lie in subtrees attached to different neighbors of c.

The process for each c is as follows:

For each neighbor v of c, traverse the subtree rooted at v (while not going back to c).
During this traversal, count how many nodes x satisfy the condition distance(c, x) % signalSpeed == 0. Note that distance(c, x) = weight(c, v) + distance(v, x).
Let's say the counts for the branches are c_1, c_2, ..., c_k. The total number of pairs is the sum of c_i * c_j for all i < j. This can be computed efficiently in a single pass over the counts. We maintain a prefix_sum of counts seen so far. For the current count c_i, we add c_i * prefix_sum to the total result for c, and then update prefix_sum by adding c_i.

Here is a Java implementation of this approach:

1import java.util.ArrayList;2import java.util.List;3 4class Solution {5    public int[] countPairsOfConnectableServers(int[][] edges, int signalSpeed) {6        int n = edges.length + 1;7        List<int[]>[] adj = new ArrayList[n];8        for (int i = 0; i < n; i++) {9            adj[i] = new ArrayList<>();10        }11        for (int[] edge : edges) {12            adj[edge[0]].add(new int[]{edge[1], edge[2]});13            adj[edge[1]].add(new int[]{edge[0], edge[2]});14        }15 16        int[] result = new int[n];17        for (int i = 0; i < n; i++) {18            if (adj[i].size() < 2) {19                continue;20            }21 22            int totalPairs = 0;23            int prefixCount = 0;24 25            for (int[] neighborInfo : adj[i]) {26                int neighbor = neighborInfo[0];27                int weight = neighborInfo[1];28                29                int countInSubtree = dfs(neighbor, i, weight, signalSpeed, adj);30                31                totalPairs += countInSubtree * prefixCount;32                prefixCount += countInSubtree;33            }34            result[i] = totalPairs;35        }36        return result;37    }38 39    private int dfs(int u, int p, int currentDist, int signalSpeed, List<int[]>[] adj) {40        int count = 0;41        if (currentDist % signalSpeed == 0) {42            count++;43        }44 45        for (int[] neighborInfo : adj[u]) {46            int v = neighborInfo[0];47            int w = neighborInfo[1];48            if (v != p) {49                count += dfs(v, u, currentDist + w, signalSpeed, adj);50            }51        }52        return count;53    }54}

Complexity

Time

O(N^2). The outer loop runs `N` times. For each server `i`, we perform a traversal for each of its neighbors. The total nodes visited across all these traversals for a single `i` is `N-1`. So, the work for each `i` is O(N). The total time complexity is O(N*N).

Space

O(N). This is for the adjacency list and the recursion stack for DFS. In the worst case (a path graph), the recursion depth can be O(N).

Trade-offs

Pros

Efficient enough for the given constraints (N <= 1000).
Relatively straightforward to implement compared to more advanced tree algorithms.

Cons

The quadratic time complexity might be too slow if N were significantly larger.

Solutions

Solution

1class Solution {2private3  int signalSpeed;4private5  List<int[]>[] g;6public7  int[] countPairsOfConnectableServers(int[][] edges, int signalSpeed) {8    int n = edges.length + 1;9    g = new List[n];10    this.signalSpeed = signalSpeed;11    Arrays.setAll(g, k->new ArrayList<>());12    for (var e : edges) {13      int a = e[0], b = e[1], w = e[2];14      g[a].add(new int[]{b, w});15      g[b].add(new int[]{a, w});16    }17    int[] ans = new int[n];18    for (int a = 0; a < n; ++a) {19      int s = 0;20      for (var e : g[a]) {21        int b = e[0], w = e[1];22        int t = dfs(b, a, w);23        ans[a] += s * t;24        s += t;25      }26    }27    return ans;28  }29private30  int dfs(int a, int fa, int ws) {31    int cnt = ws % signalSpeed == 0 ? 1 : 0;32    for (var e : g[a]) {33      int b = e[0], w = e[1];34      if (b != fa) {35        cnt += dfs(b, a, ws + w);36      }37    }38    return cnt;39  }40}

1class Solution {2public:3  vector<int> countPairsOfConnectableServers(vector<vector<int>> &edges,4                                             int signalSpeed) {5    int n = edges.size() + 1;6    vector<pair<int, int>> g[n];7    for (auto &e : edges) {8      int a = e[0], b = e[1], w = e[2];9      g[a].emplace_back(b, w);10      g[b].emplace_back(a, w);11    }12    function<int(int, int, int)> dfs = [&](int a, int fa, int ws) {13      int cnt = ws % signalSpeed == 0;14      for (auto &[b, w] : g[a]) {15        if (b != fa) {16          cnt += dfs(b, a, ws + w);17        }18      }19      return cnt;20    };21    vector<int> ans(n);22    for (int a = 0; a < n; ++a) {23      int s = 0;24      for (auto &[b, w] : g[a]) {25        int t = dfs(b, a, w);26        ans[a] += s * t;27        s += t;28      }29    }30    return ans;31  }32};

1class Solution:2    def countPairsOfConnectableServers(self, edges: List[List[int]], signalSpeed: int) -> List[int]: def dfs(a: int, fa: int, ws: int) -> int: cnt = 0 if ws % signalSpeed else 1 for b, w in g[a]: if b != fa: cnt += dfs(b, a, ws + w) return cnt 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)) ans = [0] * n for a in range(n): s = 0 for b, w in g[a]: t = dfs(b, a, w) ans[a] += s * t s += t return ans

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem