59 lines
1.8 KiB
Python
59 lines
1.8 KiB
Python
MAX_CONNS = 999
|
|
|
|
class DSU: # motherfucker i havent seen or used this in so long
|
|
def __init__(self, n):
|
|
self.parent = list(range(n))
|
|
self.size = [1] * n
|
|
|
|
def find(self, u): # we got that mf path compression and everything!
|
|
root = u
|
|
while self.parent[root] != root:
|
|
root = self.parent[root]
|
|
while u != root:
|
|
nxt = self.parent[u]
|
|
self.parent[u] = root
|
|
u = nxt
|
|
return root
|
|
|
|
def union(self, u, v):
|
|
root_u = self.find(u)
|
|
root_v = self.find(v)
|
|
if root_u != root_v: # union by size -> size of component stored at root
|
|
if self.size[root_u] < self.size[root_v]:
|
|
root_u, root_v = root_v, root_u
|
|
|
|
self.parent[root_v] = root_u
|
|
self.size[root_u] += self.size[root_v]
|
|
return True
|
|
return False
|
|
|
|
|
|
def build_graph():
|
|
edge_list = []
|
|
with open('input.txt', 'r') as file:
|
|
data = file.read().strip().split('\n')
|
|
num_vertices = len(data)
|
|
for i, line in enumerate(data):
|
|
coords = line.split(',')
|
|
for j in range(i + 1, len(data)):
|
|
if i != j:
|
|
other_coords = data[j].split(',')
|
|
distance = ((int(coords[0]) - int(other_coords[0])) ** 2 + (int(coords[1]) - int(other_coords[1])) ** 2 + (int(coords[2]) - int(other_coords[2])) ** 2) ** 0.5
|
|
edge_list.append((i, j, distance, coords[0], other_coords[0]))
|
|
return edge_list, num_vertices
|
|
|
|
|
|
if __name__ == "__main__":
|
|
graph, num_vertices = build_graph()
|
|
graph.sort(key=lambda x: x[2])
|
|
dsu = DSU(num_vertices)
|
|
conns = 0
|
|
for u, v, w, u_coord, v_coord in graph:
|
|
if conns >= MAX_CONNS:
|
|
break
|
|
if dsu.union(u, v):
|
|
conns += 1
|
|
result = int(u_coord) * int(v_coord)
|
|
|
|
print(result)
|
|
|