K-Means: assignment, centroid updates and empty clusters
Posted 26 October 2026. Original preparation material, not a claim about any company’s interview requirements.
What you should be able to do
Explain the objective, trace a small example on paper, implement the core without library algorithm shortcuts, and defend boundary behavior. Prerequisites: Python, arrays, basic linear algebra and derivatives.
Intuition
K-Means alternates two easy conditional optimization problems. Fix centers and choose the nearest one for every point. Fix assignments and choose the mean of each cluster. The model favors compact Euclidean clusters; it is not a general solution for arbitrary shapes or densities.
Formula and derivation
For a fixed cluster, differentiate the squared-error sum with respect to its center: 2|C|μ−2
Paper-and-pencil example
For points [0,2,8,10] and initial centers [0,10], assignments are [0,0,1,1]. The new centers are [1,9]. Squared error changes from 8 to 4. The next assignment does not change, so these centers are a fixed point. Labels can be permuted without changing the clustering.
Implementation contract and common errors
An empty cluster has no mean. This reference keeps its previous center; alternatives include reinitializing to a distant point, but that needs a declared policy. Recompute labels after the final center update. Time O(I n k d); the simple broadcast distance implementation temporarily uses O(nkd) memory. Production code can compute distances in batches. Standardize features when Euclidean scale is inappropriate.
Complete runnable reference
The shared file contains all five algorithms, including multi-head attention. Requires Python 3.11+ and NumPy; validated with Python 3.12.14 and NumPy 2.3.5. Inputs are deterministic; no random seed is needed. No paid API or model download.
Download Python implementation and tests
python -m pip install numpy==2.3.5
python core.py
# PASS: 10 algorithm invariants (synthetic fixtures)Show full source
import numpy as np
def softmax(x):
x = np.asarray(x, dtype=float)
if x.size == 0 or not np.isfinite(x).all():
raise ValueError('nonempty finite logits required')
z = np.exp(x - x.max(axis=-1, keepdims=True))
return z / z.sum(axis=-1, keepdims=True)
def attention(q, k, v, causal=False):
if q.shape[-1] != k.shape[-1] or k.shape[-2] != v.shape[-2]:
raise ValueError('incompatible shapes')
scores = q @ k.swapaxes(-1, -2) / np.sqrt(q.shape[-1])
if causal:
if q.shape[-2] != k.shape[-2]:
raise ValueError('this teaching mask requires equal sequence lengths')
mask = np.tril(np.ones(scores.shape[-2:], dtype=bool))
scores = np.where(mask, scores, -np.inf)
z = np.exp(scores - scores.max(axis=-1, keepdims=True))
weights = z / z.sum(axis=-1, keepdims=True)
return weights @ v
def multi_head(x, wq, wk, wv, wo, heads, causal=True):
b, t, d = x.shape
if heads < 1 or d % heads: raise ValueError('heads must divide model dimension')
def split(w):
return (x @ w).reshape(b, t, heads, d // heads).transpose(0, 2, 1, 3)
y = attention(split(wq), split(wk), split(wv), causal)
return y.transpose(0, 2, 1, 3).reshape(b, t, d) @ wo
def kmeans(x, centers, max_iter=100, tol=1e-8):
x, c = np.asarray(x, float), np.array(centers, float, copy=True)
if x.ndim!=2 or c.ndim!=2 or not len(x) or not len(c) or x.shape[1]!=c.shape[1] or max_iter<1 or not np.isfinite(x).all() or not np.isfinite(c).all():
raise ValueError('invalid data or centers')
for _ in range(max_iter):
labels = ((x[:,None,:]-c[None,:,:])**2).sum(-1).argmin(-1)
new = np.array([x[labels==j].mean(0) if np.any(labels==j) else c[j] for j in range(len(c))])
shift = np.max(np.linalg.norm(new-c,axis=1))
c = new
if shift <= tol: break
labels = ((x[:,None,:]-c[None,:,:])**2).sum(-1).argmin(-1)
return c, labels, ((x-c[labels])**2).sum()
class Adam:
def __init__(self, shape, lr=.001, b1=.9, b2=.999, eps=1e-8):
self.m=np.zeros(shape); self.v=np.zeros(shape); self.t=0
self.lr,self.b1,self.b2,self.eps=lr,b1,b2,eps
def step(self, p, g):
if p.shape!=self.m.shape or g.shape!=p.shape or not np.isfinite(g).all(): raise ValueError('invalid gradient')
self.t+=1
self.m=self.b1*self.m+(1-self.b1)*g
self.v=self.b2*self.v+(1-self.b2)*g*g
mh=self.m/(1-self.b1**self.t); vh=self.v/(1-self.b2**self.t)
return p-self.lr*mh/(np.sqrt(vh)+self.eps)
def beam_search(next_log_probs, width=2, max_steps=4, eos='EOS'):
if width<1 or max_steps<1: raise ValueError('positive limits required')
beams=[((),0.,False)]
for _ in range(max_steps):
candidates=[]
for seq,score,done in beams:
if done: candidates.append((seq,score,True)); continue
for token,lp in next_log_probs(seq).items():
if not np.isfinite(lp) or lp>0: raise ValueError('finite log probabilities required')
candidates.append((seq+(token,),score+lp,token==eos))
if not candidates: raise ValueError('no candidates')
beams=sorted(candidates,key=lambda x:(-x[1],x[0]))[:width]
if all(x[2] for x in beams): break
return beams
if __name__=='__main__':
assert np.allclose(softmax([1000,1000]),[.5,.5])
assert np.allclose(softmax([1,2]),softmax([101,102]))
x=np.arange(24,dtype=float).reshape(1,3,8)/24; eye=np.eye(8)
y=multi_head(x,eye,eye,eye,eye,2); changed=x.copy(); changed[:,2]+=10
assert np.allclose(y[:,:2],multi_head(changed,eye,eye,eye,eye,2)[:,:2])
assert np.allclose(multi_head(x,eye,eye,eye,eye,1),attention(x,x,x,True))
c,l,s=kmeans([[0],[2],[8],[10]],[[0],[10]])
assert np.allclose(c,[[1],[9]]) and s==4
c,l,s=kmeans([[0],[0]],[[0],[8]]); assert np.isfinite(c).all()
opt=Adam((1,)); assert np.allclose(opt.step(np.array([1.]),np.array([2.])),[.999])
assert opt.t==1
def probs(seq):
p={'A':.6,'B':.4} if not seq else ({'EOS':.5,'X':.5} if seq[-1]=='A' else {'EOS':.9,'X':.1})
return {k:np.log(v) for k,v in p.items()}
assert beam_search(probs,max_steps=2)[0][0]==('B','EOS')
assert beam_search(probs,width=1,max_steps=2)[0][0]==('A','EOS')
print('PASS: 10 algorithm invariants (synthetic fixtures)')
Interview follow-up
Why not stop when labels equal zero?
Hint: Distinguish labels from convergence.
Answer and reasoning
Cluster zero is a valid assignment. Stop on unchanged assignments, small center movement, or an iteration budget. Report when the budget terminates before convergence.
Practice without looking
- Spend 3 minutes writing shapes, assumptions and an example.
- Spend 12 minutes writing the core implementation on paper.
- Spend 5 minutes tracing an ordinary input and one failure case.
- Type your handwritten version and run the reference invariants. Explain every mismatch before consulting the solution.
Transfer exercise: change one assumption in the implementation contract above. Write a new test that fails before your modification and passes afterward. A convincing explanation of a limitation is better than silently inventing unsupported behavior.
Sources and next lesson
Primary paper or official documentation
Softmax: stable probabilities and gradients
Attention and Multi-Head Attention: write the shapes first
K-Means: assignment, centroid updates and empty clusters