← 返回 apple 的题目列表Implement K-Means Clustering with NumPy/PyTorch (Vectorized)
类型:online_judge
Implement K-Means clustering with an efficient vectorized NumPy or PyTorch solution (avoid per-sample Python loops).
Given:
Data matrix X of shape (n, d)
Number of clusters k
max_iters
Convergence tolerance tol (stop when the maximum centroid L2 shift <= tol)
Implement a function that returns:
centroids of shape (k, d)
labels of shape (n,), with values in [0, k-1]
Requirements
Initialization: with seed, randomly sample k points from X without replacement as initial centroids.
Iterate:
Assignment: assign each point to the nearest centroid by Euclidean distance.
Update: set each centroid to the mean of assigned points.
Empty cluster handling: if a cluster receives no points in an iteration, reinitialize its centroid to the sample that is farthest from its nearest centroid (farthest-point reinit).
Use broadcasting / tensor ops (e.g., unsqueeze/keepdim) for efficiency.
Input (stdin)
First line: n d k max_iters tol seed Next n lines: d floats per line
Output (stdout)
Line 1: n integer labels Next k lines: d floats per centroid (6 decimals)
Constraints
1 <= n <= 20000, 1 <= d <= 100, 1 <= k <= min(n, 100)
1 <= max_iters <= 300, tol > 0
Example
Input:
6 2 2 100 1e-6 0
0 0
0 1
1 0
10 10
10 11
11 10
Valid output (one possibility):
0 0 0 1 1 1
0.333333 0.333333
10.333333 10.333333
Example
Input
6 2 2 100 1e-6 0
0 0
0 1
1 0
10 10
10 11
11 10
Output
0 0 0 1 1 1
0.333333 0.333333
10.333333 10.333333