|
1 | 1 | from collections import defaultdict |
| 2 | +import queue |
2 | 3 |
|
3 | 4 |
|
4 | | -def updateCost(neighbors, b, cost): |
5 | | - for i in range(len(neighbors)): |
6 | | - if neighbors[i][0] == b: |
7 | | - neighbors[i][1] = cost |
| 5 | +def updateCost(arr, key, val): |
| 6 | + return [(k, v) if v != val else (key, val) for (k, v) in arr] |
8 | 7 |
|
9 | 8 |
|
10 | 9 | def minCostRepair(n, edges, edgesToRepair): |
11 | 10 | # create edge lists associated with cost |
12 | 11 | edgeList = defaultdict(list) |
13 | 12 | for a, b in edges: |
14 | | - edgeList[a].append([b, 0]) |
15 | | - edgeList[b].append([a, 0]) |
16 | | - # print(edgeList) |
| 13 | + edgeList[a].append((0, b)) |
| 14 | + edgeList[b].append((0, a)) |
| 15 | + print(edgeList) |
17 | 16 | for a, b, cost in edgesToRepair: |
18 | | - updateCost(edgeList[a], b, cost) |
19 | | - updateCost(edgeList[b], a, cost) |
20 | | - # print(edgeList) |
| 17 | + edgeList[a] =updateCost(edgeList[a], cost, b) |
| 18 | + edgeList[b] =updateCost(edgeList[b], cost, a) |
| 19 | + print(edgeList) |
21 | 20 | # find shorted spanning graph |
22 | 21 | reached = set() |
23 | 22 | totalCost = 0 |
24 | | - queue = list() |
25 | | - queue.append([1, 0]) |
| 23 | + pq = queue.PriorityQueue() |
| 24 | + pq.put((0, 1)) |
26 | 25 | while len(reached) < n: |
27 | | - current = queue.pop(0) |
28 | | - currentNode = current[0] |
29 | | - cost = current[1] |
| 26 | + current = pq.get() |
| 27 | + currentNode = current[1] |
| 28 | + cost = current[0] |
30 | 29 | if currentNode not in reached: |
31 | 30 | reached.add(currentNode) |
32 | 31 | totalCost += cost |
33 | | - queue.extend(edgeList[currentNode]) |
34 | | - queue.sort(key=lambdapair: pair[1]) |
| 32 | + forneighborinedgeList[currentNode]: |
| 33 | + pq.put(neighbor) |
35 | 34 | return totalCost |
36 | 35 |
|
37 | 36 |
|
|
0 commit comments