KMeans на NumPy: центроиды и инерция за один шаг

Один итерационный шаг KMeans: распределение и центроиды

Что уже дано

Проверяющая система выполняет решение в среде Python 3.12 с предустановленной библиотекой numpy. Никаких глобальных переменных или скрытых файлов перед запуском решения не создаётся. Проверяющие тесты напрямую импортируют и вызывают функцию kmeans_single_step(X, centers).

Что нужно сделать

Реализуйте функцию kmeans_single_step(X, centers).

Функция принимает:

• X: вещественный двумерный массив np.ndarray (или вложенный числовой список) формы (N, p), где $N\ge 1$, $p\ge 1$ — координаты точек выборки;
• centers: вещественный двумерный массив np.ndarray (или вложенный числовой список) формы (k, p), где $k\ge 1$ — начальные координаты центроидов.

Функция должна выполнить ровно один базовый шаг алгоритма KMeans без запуска полного цикла оптимизации scikit-learn:

1. Для каждой точки $x_i\in X$ найти индекс ближайшего центра из centers по квадрату евклидова расстояния $\Vert x_i-c_j{\Vert }^2$. Гарантируется, что ничьих по расстояниям нет и каждый из $k$ центров получает хотя бы одну точку.
2. Пересчитать координаты каждого из $k$ центров как среднее арифметическое точек, отнесённых к этому центру на шаге 1.
3. Рассчитать суммарную инерцию (сумму квадратов евклидовых расстояний от точек до их обновлённых центров) для этого же зафиксированного распределения по группам: ${\sum }_{j=0}^{k-1}{\sum }_{x_i\in C_j}\Vert x_i-{\text{new\_center}}_j{\Vert }^2$.

Функция возвращает словарь со следующими ключами:

• "labels": одномерный массив np.ndarray длины $N$ с целочисленным типом (например, int64), содержащий индексы назначенных центров ($0,1,\dots ,k-1$);
• "new_centers": двумерный вещественный массив np.ndarray формы (k, p) с обновлёнными координатами центроидов;
• "inertia": вещественное число float — сумма квадратов евклидовых расстояний от точек до обновлённых центроидов при закреплённом распределении.

Ввод и вывод

stdin не используется. Печатать через print в stdout ничего не нужно. Проверяющая система вызывает kmeans_single_step и передаёт аргументы в виде числовых массивов или списков.

Пример использования

import numpy as np

X = np.array([[1.0, 0.0], [3.0, 0.0], [5.0, 0.0]])
centers = np.array([[0.0, 0.0], [4.0, 0.0]])

result = kmeans_single_step(X, centers)
print(result["labels"].tolist())
print(np.round(result["new_centers"], 4).tolist())
print(round(result["inertia"], 4))

Результат:

[0, 1, 1]
[[1.0, 0.0], [4.0, 0.0]]
2.0