手撕-Kmeans

import numpy as np
from typing import List, Tuple

def kmeans(n: int, k: int, max_iter: int, nums: List[List[float]]) -> Tuple[List[List[str]], List[int]]:
    points = np.array(nums, dtype=float)

    # 1. 初始化中心（取前k个点）
    centers = points[:k].copy()
    # 初始化标签
    labels = np.zeros(n, dtype=int)
    

    for _ in range(max_iter):
        new_labels = []
        
        # 2. 寻找最近簇中心
        for point in points:
            # [k, 2] - [2] -> [k,2] -> [k]
            dists = np.sum((centers - point) ** 2, axis=1)
            closest_center = np.argmin(dists)
            new_labels.append(closest_center)
        
        # 这里不要忘记转类型
        new_labels = np.array(new_labels)
        
        # 3. 判断是否收敛：标签不再变化
        if np.array_equal(labels, new_labels):
            break
        
        labels = new_labels
        
        # 4. 更新簇中心坐标：属于该簇的所有坐标取平均值
        for j in range(k):
            cluster_points = points[labels == j]
            if len(cluster_points) > 0:
                centers[j] = np.mean(cluster_points, axis=0)
    

    centers_str = [[f"{x:.4f}", f"{y:.4f}"] for x, y in centers]
        
    return centers_str, labels.tolist()
    
    
if __name__ == "__main__":
    nums = [
        [1.0, 1.0], [1.5, 1.5], [1.2, 1.8],  # 簇 0
        [5.0, 5.0], [5.5, 5.5], [4.8, 5.2]   # 簇 1
    ]

    n = len(nums)    
    k = 2
    max_iter = 10
    
    centers, labels = kmeans(n, k, max_iter, nums)
    # 簇中心
    print(centers)
    # 标签列表
    print(labels)