114 lines
2.4 KiB
Python
114 lines
2.4 KiB
Python
import math
|
|
from math import floor, sqrt
|
|
|
|
def dist(p,q):
|
|
return sqrt(sum((x-y)**2 for x,y in zip(p,q)))
|
|
|
|
def solve(input, limit):
|
|
points = []
|
|
for line in open(input):
|
|
x,y,z = map(int,line.split(','))
|
|
points.append((x,y,z))
|
|
|
|
def nearest(p):
|
|
mindist = float('inf')
|
|
minpoint = p
|
|
for q in points:
|
|
if q != p:
|
|
d = dist(p,q)
|
|
if d < mindist:
|
|
mindist = d
|
|
minpoint = q
|
|
return minpoint
|
|
|
|
#for p in points:
|
|
# q = shrink(p)
|
|
# print(p, nearest(p), list(chain(*[m.get(x,[]) for x in near(q)])))
|
|
|
|
pairs = []
|
|
for p in points:
|
|
for q in points:
|
|
#q = nearest(p)
|
|
if p != q:
|
|
d = dist(p,q)
|
|
pairs.append((d,p,q))
|
|
|
|
pairs.sort()
|
|
print(len(pairs))
|
|
print(*pairs[:10], sep='\n')
|
|
|
|
circuit = []
|
|
direct = set()
|
|
parent = {p:p for p in points}
|
|
size = {p:1 for p in points}
|
|
|
|
def find(x):
|
|
root = x
|
|
while parent[root] != root:
|
|
root = parent[root]
|
|
while parent[x] != root:
|
|
x, parent[x] = parent[x], root
|
|
return root
|
|
|
|
def union(p,q):
|
|
p = find(p)
|
|
q = find(q)
|
|
if p == q:
|
|
return
|
|
if size[p] < size[q]:
|
|
p,q = q,p
|
|
parent[q] = p
|
|
size[p] += size[q]
|
|
|
|
# part 1
|
|
|
|
n = 0
|
|
for _,p,q in pairs:
|
|
if n >= limit:
|
|
break
|
|
if (p,q) in direct or (q,p) in direct:
|
|
# already directly connected
|
|
continue
|
|
print("connecting", p, q)
|
|
n += 1
|
|
union(p,q)
|
|
direct.add((p,q))
|
|
|
|
seen = set()
|
|
sizes = []
|
|
for p in points:
|
|
r = find(p)
|
|
if r in seen:
|
|
continue
|
|
seen.add(r)
|
|
print(r, size[r])
|
|
sizes.append(size[r])
|
|
|
|
sizes.sort(reverse=True)
|
|
print(sizes[:3])
|
|
|
|
t = 1
|
|
for x in sizes[:3]:
|
|
t *= x
|
|
print(t)
|
|
|
|
# part 2 (continued)
|
|
for _,p,q in pairs:
|
|
if (p,q) in direct or (q,p) in direct:
|
|
# already directly connected
|
|
continue
|
|
print("connecting", p, q)
|
|
n += 1
|
|
union(p,q)
|
|
direct.add((p,q))
|
|
r = find(p)
|
|
if size[r] == len(points):
|
|
print(p,q)
|
|
print(p[0]*q[0])
|
|
break
|
|
|
|
|
|
solve("sample", 10)
|
|
solve("input", 1000)
|
|
|