def k_means(points, init_centeroids):
"""LLoyd's algorithm on 3d-points."""
dist = INF
(x0, y0, z0), (x1, y1, z1) = init_centeroids
for _ in range(MAX_ITER):
prev_dist, dist = dist, 0.0
cluster0, cluster1 = [0, 0, 0], [0, 0, 0]
count0 = count1 = 0
for x, y, z in points:
d0 = (x0 - x) * (x0 - x) + (y0 - y) * (y0 - y) + (z0 - z) * (z0 - z)
d1 = (x1 - x) * (x1 - x) + (y1 - y) * (y1 - y) + (z1 - z) * (z1 - z)
if d0 < d1:
dist += d0
count0 += 1
c = cluster0
else:
dist += d1
count1 += 1
c = cluster1
c[0] += x
c[1] += y
c[2] += z
if abs(dist - prev_dist) < EPS:
break
x0, y0, z0 = [x / count0 for x in cluster0]
x1, y1, z1 = [x / count1 for x in cluster1]
return dist