[Python problem solution] 100226. Count the number of connectable server pairs in a weighted tree network

Topic: 100226. Counting the number of connectable server pairs in a weighted tree network

Question description

You are given an unrooted weighted tree with n vertices, representing from 0 to n - 1 Numbered servers, an array edges, among which edges[i] = [ai, bi, weighti] represents the vertex ai and bi The bidirectional edge between them, the edge weight is weighti. You are also given an integer signalSpeed。

If the following conditions are met, two servers a and b via server c Connection:

• a < b，a != c and b != c。
• from c Arrive a The distance can be signalSpeed Divisible.
• from c Arrive b The distance can be signalSpeed Divisible.
• from c Arrive b He Cong c Arrive a The paths do not share any edges.

Returns an array of integers count, the length is n, among which count[i] is through the server i Number of connectable server pairs.

Example

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: due to signalSpeed is 1,count[c] equal to from c The number of pairs of paths that originate and do not share any edges.
In a given path diagram,count[c] equal to c Number of servers on the left multiplied by c Number of servers on the right.

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 proven that servers other than 0 and 6 cannot connect to any two servers.

Restrictions

• 2 <= n <= 1000
• edges.length == n - 1
• edges[i].length == 3
• 0 <= ai, bi < n
• 1 <= weighti <= 10^6
• 1 <= signalSpeed <= 10^6
• input guarantee edges Represents a valid tree.

Problem overview

Given a tree structure representing a network of servers. Each server is connected to other servers through weighted edges. The goal is to count the number of connectable server pairs through each server in the tree, these conditions are given bysignalSpeeddefinition.

Conditions for connecting to the server

If the following conditions are met, two servers a and b via server c Connection:

• a < b
• a and b are different from c。
• from c Arrive a He Cong c Arrive b distance can be signalSpeed Divisible.
• from c Arrive a He Cong c Arrive b The paths do not share any edges.

Solution

1. tree representation:The tree is represented using an adjacency list, with each server connected to its neighbors and the weights of the connecting edges.

2. Depth First Search (DFS)：

   • The tree is traversed starting from each server using a modified DFS algorithm.
   • This algorithm calculates the distance of all other servers from the current server.
   • DFS ensures that shared edges are not considered in the path.

3. Count connectable pairs：

   • for each server c, the algorithm identifies by removing c All possible subtrees formed.
   • It computes in each subtree the c distance can be signalSpeed Divisible by the number of servers.
   • via server c The total number of connectable pairs of is calculated by considering all possible combinations of pairs from different subtrees.

Python implementation

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

from collections import defaultdict

def count_connectable_servers(edges, signal_speed):
    def build_tree(edges):
        tree = defaultdict(list)
        for a, b, weight in edges:
            tree[a].append((b, weight))
            tree[b].append((a, weight))
        return tree

    def dfs_count_nodes(server, parent, distance):
        if distance % signal_speed == 0:
            count[0] += 1
        for neighbor, weight in tree[server]:
            if neighbor != parent:
                dfs_count_nodes(neighbor, server, distance + weight)

    n = len(edges) + 1
    tree = build_tree(edges)
    counts = [0] * n

    for c in range(n):
        subtree_counts = []
        for neighbor, weight in tree[c]:
            count = [0]
            dfs_count_nodes(neighbor, c, weight)
            subtree_counts.append(count[0])

        for i in range(len(subtree_counts)):
            for j in range(i + 1, len(subtree_counts)):
                counts[c] += subtree_counts[i] * subtree_counts[j]

    return counts

Complexity analysis

• time complexity: This algorithm performs DFS once for each server in the tree. Since each edge is visited once in each DFS, and there are n servers, so the overall time complexity is O(n^2), where n is the number of servers.
• space complexity: The space complexity is O(n) due to the storage tree structure and auxiliary data structures used in the DFS process.

test case

The solution has been tested with the provided examples and additional custom test cases to ensure its correctness.

Conclusion

This solution effectively counts the number of connectable server pairs through each server in a tree network, taking into account signalSpeed Given restrictions associated with connection rules.