np_kmeans.py en Python

Agrupar puntos 2-D con el k-means de Lloyd en NumPy.

"""Agrupa puntos 2-D con unas rondas del algoritmo de Lloyd."""
import numpy as np


def kmedias(puntos, k, iteraciones=10, semilla=0):
    rng = np.random.default_rng(semilla)
    centros = puntos[rng.choice(len(puntos), size=k, replace=False)]
    etiquetas = np.zeros(len(puntos), dtype=int)
    for _ in range(iteraciones):
        distancias = np.linalg.norm(
            puntos[:, np.newaxis] - centros[np.newaxis, :], axis=2)
        etiquetas = distancias.argmin(axis=1)
        for c in range(k):
            miembros = puntos[etiquetas == c]
            if len(miembros):
                centros[c] = miembros.mean(axis=0)
    return centros, etiquetas


def principal():
    rng = np.random.default_rng(1)
    nube_a = rng.normal(loc=[0, 0], scale=0.5, size=(50, 2))
    nube_b = rng.normal(loc=[5, 5], scale=0.5, size=(50, 2))
    puntos = np.vstack([nube_a, nube_b])
    centros, etiquetas = kmedias(puntos, k=2)
    for c, centro in enumerate(centros):
        size = int((etiquetas == c).sum())
        print(f"grupo {c}: {size} puntos cerca de {centro.round(2)}")


if __name__ == "__main__":
    principal()

Cómo funciona

1. El broadcasting calcula cada distancia punto-centro.
2. argmin asigna cada punto a su centro más cercano.
3. Los centros se mueven a la media de sus miembros en cada ronda.

Palabras clave y builtins usados aquí

• as
• def
• enumerate
• for
• if
• int
• len
• print
• range
• return

El intento, en números

Líneas

32

Caracteres a escribir

984

Tokens

303

Ritmo de tres estrellas

115 tpm

Al ritmo de tres estrellas de 115 tokens por minuto, este intento toma unos 158 segundos.

Escribe este fragmento

Paso 1 de 1 en Bis; paso 22 de 22 en NumPy.

← Anterior