-
Notifications
You must be signed in to change notification settings - Fork 21
Expand file tree
/
Copy pathmatrix.py
More file actions
98 lines (78 loc) · 2.93 KB
/
matrix.py
File metadata and controls
98 lines (78 loc) · 2.93 KB
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
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
#!/bin/python3
import math
import os
import random
import re
import sys
# Let's use an adapted form of Kruskal.
# 1. Construct a set for each city.
# 2. Iterate through edges in descending order of "cost to destroy".
# 3. If the edge connects two machines, we must destroy it so add to total time.
# If the edge does not connect two machines, we do not need to destroy it.
# 4. After iterating through all edges, return total time
# From https://en.wikipedia.org/wiki/Kruskal%27s_algorithm#Pseudocode
# KRUSKAL(G):
# A = ∅
# foreach v ∈ G.V:
# MAKE-SET(v)
# foreach (u, v) in G.E ordered by weight(u, v), increasing:
# if FIND-SET(u) ≠ FIND-SET(v):
# A = A ∪ {(u, v)}
# UNION(u, v)
# return A
class Matrix:
MACHINE = "MACHINE"
def __init__(self, roads, num_cities, machines):
self.roads = roads
self.num_cities = num_cities
self.machines = machines
self.sets = [[city] for city in range(num_cities)]
self.set_member_lookup = {city: city for city in range(num_cities)}
self.mark_cities_with_machines(machines)
def mark_cities_with_machines(self, machines):
for machine in machines:
city_with_machine = self.get_set_for_city(machine)
city_with_machine.append(self.MACHINE)
def get_set_for_city(self, city):
return self.sets[self.find_set(city)]
def find_set(self, city):
return self.set_member_lookup[city]
def build_road(self, city_a, city_b):
set_a = self.get_set_for_city(city_a)
set_b = self.get_set_for_city(city_b)
if set_a is not None and set_b is not None:
set_a.extend(set_b)
set_a_lookup = self.find_set(city_a)
for city in set_b:
self.set_member_lookup[city] = set_a_lookup
set_b = None
def min_time(self):
"""
Return an integer representing the minimum time required to disrupt the
connections among all machines.
"""
total_time = 0
descending_roads = list(
reversed(
sorted(
self.roads,
key=lambda tuple: tuple[2])))
for road in descending_roads:
city_a, city_b, time_to_destroy = road
set_a = self.get_set_for_city(city_a)
set_b = self.get_set_for_city(city_b)
if self.MACHINE in set_a and self.MACHINE in set_b:
total_time += time_to_destroy
else:
self.build_road(city_a, city_b)
return total_time
if __name__ == '__main__':
num_cities, num_machines = map(int, input().split())
roads = []
for _ in range(num_cities - 1):
roads.append(tuple(map(int, input().rstrip().split())))
machines = [] # these are cities machines are in
for _ in range(num_machines):
machines.append(int(input()))
matrix = Matrix(roads, num_cities, machines)
print(matrix.min_time())