One step of cosine k-means

Vector indexes like IVF split the space into clusters, then search only the clusters closest to the query. The clusters come from k-means. Each round has two halves: assign every vector to a centroid, then move every centroid to the middle of its vectors. Write one round.

Write kmeans_step(vectors, centroids):

• vectors and centroids are lists of float lists, all the same length. There is at least one centroid.
• Assign: each vector goes to the centroid with the highest cosine similarity. Break ties by the lower centroid index. A zero vector has similarity 0 with everything.
• Update: each new centroid is the element-wise mean of the vectors assigned to it. A centroid with no vectors keeps its old position.
• Return {"assignments": [...], "centroids": [...]}: one centroid index per vector, and the new centroids as lists of floats.

No numpy. Plain Python only.

>>> kmeans_step([[1, 0], [3, 1], [0, 2]], [[1, 0], [0, 1]])
{'assignments': [0, 0, 1], 'centroids': [[2.0, 0.5], [0.0, 2.0]]}

Approach

Two passes. Score every vector against every centroid and keep the best. Then sum the members of each cluster and divide.

import math


def kmeans_step(vectors, centroids):
    def norm(v):
        return math.sqrt(sum(x * x for x in v))

    k = len(centroids)
    dim = len(centroids[0])
    cnorms = [norm(c) for c in centroids]

    assignments = []
    for v in vectors:
        vn = norm(v)
        best, best_score = 0, None
        for j, c in enumerate(centroids):
            if vn == 0 or cnorms[j] == 0:
                s = 0.0
            else:
                s = sum(a * b for a, b in zip(v, c)) / (vn * cnorms[j])
            if best_score is None or s > best_score:
                best, best_score = j, s
        assignments.append(best)

    sums = [[0.0] * dim for _ in range(k)]
    counts = [0] * k
    for v, j in zip(vectors, assignments):
        counts[j] += 1
        row = sums[j]
        for d, x in enumerate(v):
            row[d] += x

    new = []
    for j in range(k):
        if counts[j] == 0:
            new.append([float(x) for x in centroids[j]])
        else:
            new.append([x / counts[j] for x in sums[j]])
    return {"assignments": assignments, "centroids": new}

The strict s > best_score keeps the first centroid that reached the best score, which is the lower index. Centroid lengths are computed once, outside the loop.

Complexity

Assignment is O(n·k·d) for n vectors, k centroids and dimension d. The update is O(n·d). The assignment dominates, which is why large-scale k-means samples the data or uses a smarter search for the nearest centroid.

Why this matters

An IVF index is this function, run until the centroids stop moving. At query time, the index compares the query with the k centroids first, then scans only the few closest clusters. That turns a scan of a million vectors into a scan of a few thousand. Strictly, cosine k-means (spherical k-means) normalises vectors before averaging. Many systems normalise every embedding up front, and then cosine and dot product give the same ranking. Empty clusters are a real problem at scale. Keeping the old centroid is the simplest fix; re-seeding it from a far-away point is the common better one.