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
34
35
36
37
38
39
40
41
42
| class Solution {
public int[] countPairsOfConnectableServers(int[][] edges, int signalSpeed) {
int n = edges.length + 1;
List<int[]>[] nodeToNodeListMap = new ArrayList[n];
for (int i = 0; i < nodeToNodeListMap.length; i++) {
nodeToNodeListMap[i] = new ArrayList<>();
}
for (int[] edge : edges) {
int nodeA = edge[0], nodeB = edge[1], weight = edge[2];
nodeToNodeListMap[nodeA].add(new int[]{nodeB, weight});
nodeToNodeListMap[nodeB].add(new int[]{nodeA, weight});
}
int[] res = new int[n];
for (int node = 0; node < res.length; node++) {
int preSum = 0;
for (int[] to : nodeToNodeListMap[node]) {
int nextNode = to[0], weight = to[1];
int count = dfs(nextNode, node, weight, nodeToNodeListMap, signalSpeed);
res[node] += preSum * count;
preSum += count;
}
}
return res;
}
private int dfs(int node, int from, int weights, List<int[]>[] nodeToNodeListMap, int signalSpeed) {
int count = weights % signalSpeed == 0 ? 1 : 0;
for (int[] to : nodeToNodeListMap[node]) {
int nextNode = to[0], weight = to[1];
if (nextNode != from) {
count += dfs(nextNode, node, weights + weight, nodeToNodeListMap, signalSpeed);
}
}
return count;
}
}
|