1 files changed, 143 insertions, 0 deletions
diff --git a/kmeans.py b/kmeans.py
new file mode 100755
index 0000000..a88e773
--- /dev/null
+++ b/kmeans.py
@@ -0,0 +1,143 @@
+#!/usr/bin/env python3
+"""Simple k-means implementation"""
+import argparse
+import math
+import random
+
+
+def distance(p1, p2) -> float:
+    """return the euclidean distance between to points"""
+
+    assert (len(p1) == len(p2))
+
+    if isinstance(p1, tuple):
+        p1 = list(p1)
+    if isinstance(p2, tuple):
+        p2 = list(p2)
+
+    d = 0
+    for i, x in enumerate(p1):
+        d += (x - p2[i])**2
+
+    return math.sqrt(d)
+
+
+def gen_points(n, dimensions, max_value, min_value) -> list:
+    """generate random points"""
+    points = []
+    for _ in range(n):
+        point = []
+        for _ in range(dimensions):
+            point.append(random.randint(min_value, max_value))
+        points.append(point)
+    return points
+
+
+def init_clusters(n, points) -> list:
+    """create random initial clusters
+    
+    The clusters are chosen randomly between the encountered maximum and
+    minimum values in the given points"""
+
+    max_values = []
+    min_values = []
+
+    for i in range(len(points[0])):
+        max_values.append(max(p[i] for p in points))
+        min_values.append(min(p[i] for p in points))
+
+    clusters = []
+    for _ in range(n):
+        cluster = []
+        for i, max_value in enumerate(max_values):
+            cluster.append(random.randint(min_values[i], max_value))
+        clusters.append(tuple(cluster))
+
+    return clusters
+
+
+def map_points(clusters, points) -> dict:
+    """map points to clusters
+
+    Map each points to the cluster center with the minimal distance."""
+
+    mapping = {p: [] for p in clusters}  # type: ignore
+    for p in points:
+        min_distance = None
+        center = None
+        for cluster in clusters:
+            d = distance(p, cluster)
+            if not min_distance or d < min_distance:
+                center = cluster
+                min_distance = d
+        mapping[center].append(p)
+
+    return mapping
+
+
+def calc_clusters(mapping) -> list:
+    """Calculate and return new centers for each cluster"""
+    # return the new centers, use the old ones if no points are in the cluster
+    return [
+        k_mean(points) if len(points) else center
+        for center, points in mapping.items()
+    ]
+
+
+def k_mean(points) -> tuple:
+    """return the k-means of the given points"""
+    # Start with the values of the first points
+    means = list(points[0])
+
+    # Only consider the remaining points
+    for p in points[1:]:
+        # calculate a running mean for each value
+        for i, x in enumerate(p):
+            means[i] = (means[i] + x) / 2
+
+    return tuple(means)
+
+
+def visualize(mapping):
+    """visualize a mapping using matplotlib"""
+    import matplotlib.pyplot as plt
+
+    for center, points in mapping.items():
+        x = [p[0] for p in points]
+        y = [p[1] for p in points]
+        plt.scatter(x, y, alpha=0.5)
+        plt.plot(center[0], center[1], 'x')
+    plt.show()
+
+
+def main():
+    """main entry point"""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-n", "--clusters", type=int, default=3)
+    parser.add_argument("-r", "--random-points", type=int, default=10)
+    parser.add_argument("-d", "--point-dimensions", type=int, default=2)
+    parser.add_argument("--dont-visualize", action='store_true')
+    parser.add_argument("--max", type=int, default=10)
+    parser.add_argument("--min", type=int, default=1)
+    args = parser.parse_args()
+
+    points = gen_points(args.random_points, args.point_dimensions, args.max,
+                        args.min)
+    clusters = init_clusters(args.clusters, points)
+    while True:
+        mapping = map_points(clusters, points)
+
+        if not args.dont_visualize:
+            visualize(mapping)
+
+        new_clusters = calc_clusters(mapping)
+
+        if new_clusters == clusters:
+            break
+        clusters = new_clusters
+
+    print(clusters)
+
+
+if __name__ == '__main__':
+    main()