Untitled

import math
import random
import matplotlib.pyplot as plt

class KMeans:
    def __init__(self, k, max_iterations=10):
        """Inicjalizacja danych algorytmu k-średnich"""

        self.k = k
        self.max_iterations = max_iterations
        self.centroids = []
        self.clusters = []

    def _distance(self, point_1, point_2):
        """Metoda do liczenie długości Euklidesowej"""

        return math.sqrt((point_1[0] - point_2[0]) ** 2 + (point_1[1] - point_2[1]) ** 2)

    def _init_centroids(self, data):
        """Metoda wybierające losowe środki początkowe"""

        self.centroids = random.sample(data, self.k)

    def _assign_clusters(self, data):
        """Metoda do przypisywania punktów do wybranych klastrów (grup)"""

        clusters = [[] for _ in range(self.k)]

        for point in data:
            distances = [self._distance(point, c) for c in self.centroids]
            cluster_id = distances.index(min(distances))
            clusters[cluster_id].append(point)

        return clusters

    def _update_centroids(self, clusters, data):
        """Metoda do aktualizacji środków"""

        new_centroids = []

        for cluster in clusters:
            if not cluster:
                new_centroids.append(random.choice(data))
                continue

            mean_x = sum(p[0] for p in cluster) / len(cluster)
            mean_y = sum(p[1] for p in cluster) / len(cluster)
            new_centroids.append((mean_x, mean_y))

        return new_centroids

    def fit(self, data):
        """Uruchomienie algorytmu"""

        self._init_centroids(data)

        for _ in range(self.max_iterations):
            self.clusters = self._assign_clusters(data)
            self.centroids = self._update_centroids(self.clusters, data)

        return self

    def plot_figure(self, filename="k-means.png"):
        """Metoda do rysowania i zapisywania wykresu."""

        colors = ['red', 'blue', 'green', 'purple']

        for i, cluster in enumerate(self.clusters):
            xs = [p[0] for p in cluster]
            ys = [p[1] for p in cluster]
            plt.scatter(xs, ys, color=colors[i], marker='o', label=f'Cluster {i + 1}')

        cx = [c[0] for c in self.centroids]
        cy = [c[1] for c in self.centroids]
        plt.scatter(cx, cy, color='black', marker='x', s=150, label='Centroids')

        plt.xlabel('x')
        plt.ylabel('y')
        plt.legend()
        plt.title('K-średnich')

        # Zapis do pliku
        plt.savefig("output/" + filename)

        plt.show()