We use cookies (including Google cookies) to personalize ads and analyze traffic. By continuing to use our site, you accept our Privacy Policy.

SimplyLeet
TagsCompaniesDisclaimersPrivacyTermsContactAbout Us

Count Pairs of Connectable Servers in a Weighted Tree Network

Number: 3326

Difficulty: Medium

Paid? No

Companies: thoughtspot, UBS

Problem Description

You are given an unrooted weighted tree with n vertices (servers numbered from 0 to n – 1) and an array edges where each edge is represented as [ai, bi, weighti]. Additionally, you are given an integer signalSpeed. Two servers a and b (with a < b) are connectable through a server c (where a ≠ c and b ≠ c) if both of the following hold:

  1. The distance from c to a is divisible by signalSpeed.
  2. The distance from c to b is divisible by signalSpeed.
  3. The two paths from c to a and c to b do not share any edges. The goal is to return an integer array count of length n where count[i] is the number of server pairs that are connectable through server i.

Key Insights

  • For each server c, consider each branch (i.e. each subtree obtained by removing c from the tree).
  • Within each branch, use DFS to compute the accumulated distance modulo signalSpeed from c. We are only interested in vertices where the distance mod signalSpeed is 0.
  • Since the connectable servers must lie in two distinct branches of c (to ensure the paths do not share any edge), count the number of valid vertices (distance mod==0) in each branch and then sum over all distinct pairs of branches.
  • The final answer for server c is the sum over i < j of (valid_count in branch i * valid_count in branch j).

Space and Time Complexity

Time Complexity: O(n^2) in the worst-case, since for each server as candidate center we perform DFS in its branches and n ≤ 1000. Space Complexity: O(n) for recursion and adjacency list storage.

Solution

For each server c, we treat it as the potential center through which connections could be made. For every neighbor of c, perform a DFS (avoiding going back to c) to calculate the cumulative distance from c. Count nodes whose cumulative distance modulo signalSpeed is 0. Because any two nodes that are connectable through c must lie in different branches (ensuring non-overlap of edge paths), the number of connectable pairs for server c is computed by taking the sum over different branch pairs (multiply the counts from the branches). Finally, aggregate the results into a count array.

Data structures used:

  • Adjacency list to represent the tree.
  • Recursion (DFS) to traverse each branch.
  • A list or counter to keep track of valid nodes (meeting the modulo condition) in each branch.

Algorithmic approach:

  1. Build the adjacency list.
  2. For each server c, iterate all its neighbors and perform a DFS from that neighbor with c as the parent, calculating the cumulative distance.
  3. In DFS, if the computed distance modulo signalSpeed equals 0, consider this node valid.
  4. Maintain a list of numbers, each representing the count of valid nodes from a separate branch.
  5. Use pairwise multiplication among these branch counts to compute the number of connectable pairs through c.
  6. Return the final count array.

Key gotcha: Ensuring that the DFS does not walk back to the center server c (acting as the parent), and correctly computing the modulo distance for each branch.

Code Solutions

# Python solution with detailed line-by-line comments
def countConnectablePairs(edges, signalSpeed):
    from collections import defaultdict

    n = len(edges) + 1  # n vertices for n-1 edges
    # Build the graph as an adjacency list
    graph = defaultdict(list)
    for u, v, w in edges:
        graph[u].append((v, w))
        graph[v].append((u, w))
    
    # Helper DFS function to count valid nodes in a branch
    # curr: current node, parent: node from which we reached curr, cum_dist: cumulative distance from the center
    def dfs(curr, parent, cum_dist):
        count_valid = 0
        # If the current distance mod signalSpeed is 0, count this node as valid
        if cum_dist % signalSpeed == 0:
            count_valid += 1
        # Visit all children except the parent
        for neighbor, weight in graph[curr]:
            if neighbor == parent:
                continue
            # Pass updated cumulative distance to DFS call
            count_valid += dfs(neighbor, curr, cum_dist + weight)
        return count_valid

    # Initialize the resulting count array with zeros
    result = [0] * n
    
    # For each server taken as center, compute the number of connectable pairs
    for center in range(n):
        branch_valid_counts = []
        # For each neighbor, perform DFS to get valid node count in that branch
        for neighbor, weight in graph[center]:
            # Start DFS from neighbor while avoiding the center node
            valid_count = dfs(neighbor, center, weight)
            # Append count if any valid node is found in the branch
            branch_valid_counts.append(valid_count)
        
        # For every pair of branches (i < j), add the product of valid counts
        pairs = 0
        m = len(branch_valid_counts)
        for i in range(m):
            for j in range(i+1, m):
                pairs += branch_valid_counts[i] * branch_valid_counts[j]
        result[center] = pairs
    return result

# Example usage:
edges1 = [[0,1,1],[1,2,5],[2,3,13],[3,4,9],[4,5,2]]
signalSpeed1 = 1
print(countConnectablePairs(edges1, signalSpeed1))  # Expected Output: [0,4,6,6,4,0]
← Back to All Questions