84 lines
2.5 KiB
Python
84 lines
2.5 KiB
Python
DESIRED_CONNS = 1000
|
|
|
|
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!
|
|
# if self.parent[u] == u:
|
|
# return u
|
|
# root = self.find(self.parent[u])
|
|
# self.parent[u] = root # we got that mf path compression and everything!
|
|
# return root
|
|
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))
|
|
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
|
|
attempted = 0
|
|
for u, v, w in graph:
|
|
attempted += 1
|
|
if dsu.union(u, v):
|
|
if attempted >= 1000:
|
|
break
|
|
|
|
print(f"Attempted edges: {attempted}, Connections made: {conns}")
|
|
component_sizes = {}
|
|
for i in range(num_vertices):
|
|
root = dsu.find(i)
|
|
size_of_component = dsu.size[root]
|
|
if root not in component_sizes:
|
|
component_sizes[root] = size_of_component
|
|
|
|
largest_components = sorted(component_sizes.values(), reverse=True)[:3]
|
|
result = 1
|
|
for size in largest_components:
|
|
result *= size
|
|
print(result)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|