Kazakhstan FAIO 2024 · Correction by IOAI Training

D. K-means

This page gives the full approach. Read it once you have searched.

Metric: m = 1 - L/L_baseline vs reference
Evaluated on: self-generated tests (official tests are hidden)
Our score: 300/300 self-tests pass the 0.95 rule
Expected level: m >= 0.95 x reference m on each test

Expected level reached

We solved the problem ourselves: downloaded the data, ran the experiments and measured the score below. The statement and data belong to their organisers.

The problem in brief

You are given up to a thousand points in the plane and $k$ starting positions for cluster centres (between 2 and 9 of them). The program must run the k-means algorithm from those starting centres and print where the centres end up. The input and output go through standard input and output. The statement fixes the initial centroids precisely so that everybody's algorithm follows the same path.

How you are scored

The output is not compared character by character. The judge measures the quality of your centroids with the k-means objective, the mean squared distance from each point to its nearest centroid:

L = \frac{1}{n}\sum_{i=1}^{n} \min_{j} \lVert \vec r_i - \vec\mu_j \rVert^2 .

It compares $L$ with the baseline $L_{\text{baseline}}$ of a single centroid at the mean of all points (the $k=1$ solution):

m = \max\!\Big(0,\; 1 - \frac{L}{L_{\text{baseline}}}\Big).

A test is worth 100 points if your $m$ is at least $0.95$ times the reference solution's $m$ , and 0 otherwise. Two consequences: the order and the formatting of the centroids do not matter, and a better clustering than the reference is always accepted. The official tests are hidden, so we validate on the sample and on 300 random tests of our own.

The key idea

Lloyd's algorithm alternates two steps, and each one can only lower $L$ .

Assignment. With the centroids fixed, give each point to its nearest centroid. For each point, this picks the smallest of the $k$ squared distances, so no other assignment gives a smaller $L$ .
Update. With the clusters fixed, move each centroid to the mean of its points. For a cluster $S$ , the function $\mu \mapsto \sum_{r\in S}\lVert r-\mu\rVert^2$ has gradient $-2\sum_{r\in S}(r-\mu)$ , which is zero exactly at the mean, $\mu = \frac{1}{|S|}\sum_{r\in S} r$ . So the mean is the best possible centre for that cluster.

$L$ never increases, and there are only finitely many ways to split $n$ points into $k$ clusters, so the algorithm must stop. We stop as soon as an assignment step changes nothing; then the next update would not move anything either.

What Lloyd's algorithm does not guarantee is the global minimum. Finding the best clustering is NP-hard, and Lloyd only finds a local minimum, which depends on where it starts:

Two centroids share the bottom-left blob, one sits between the two top blobs: Lloyd converges in 4 iterations to L = 125, while the best clustering has L = 17

In this example, two centroids start in the same blob and one sits halfway between two blobs. Each step is locally optimal, so nothing ever pulls a centroid across the empty space, and the algorithm stalls at $L = 125$ . The good clustering has $L = 17$ .

Two details that matter

Empty clusters. The initial centroids are arbitrary points in $(1, 50)^2$ , so a centroid may get no points at all. Its mean is then undefined. We leave it where it is. Libraries differ here: scikit-learn, for instance, moves an empty centroid to a far-away point. On our 300 random tests, our plain Lloyd matches scikit-learn's Lloyd on all 103 tests with no empty cluster, but on only 6 of the 197 tests where a cluster empties. We cannot know which convention the reference uses.

Insurance, thanks to the metric. Since a lower $L$ can only raise $m$ , we run Lloyd from the given centroids, which is what the reference does, plus 5 restarts from k-means++ initialisations (each new centre is drawn with probability proportional to its squared distance to the centres already chosen). We print the run with the lowest $L$ . Whatever the reference does with empty clusters or ties, our $L$ is at most that of the given-start Lloyd, and usually lower.

The full solution

def lloyd(points, cents):
    labels = None
    while True:
        new = [argmin_j(dist2(p, cents[j])) for p in points]   # assignment
        if new == labels:
            return cents                                       # converged
        labels = new
        cents = [mean(points in cluster j) or cents[j] for j]  # update (empty: stay)

best = min([lloyd(pts, init)] + [lloyd(pts, kmeanspp(pts, k)) for _ in range(5)], key=L)

Parsing reads all tokens with split(), so the blank lines of the sample are harmless. Each iteration costs $O(nk)$ , which is tiny here ( $n<1000$ , $k<10$ ).

On the sample, the three points on the left go to $(-2, 2)$ and the two on the right to $(3, 0)$ , so the centroids become $(-4/3,\ 4/3)$ and $(1.5,\ 0)$ , and the next assignment changes nothing. The statement's printed answer, "1.5 0" and "-1.5 1.3", is not the mean of any group of the points: it looks like a rounded or mistyped version of $(-1.333, 1.333)$ . Our centroids give a lower loss ( $L = 0.367$ against $0.384$ ), so they score at least as well under the metric.

Results

check (300 random tests, $n$ up to 999, $k$ from 2 to 9)	result
sample	$(-1.333, 1.333)$ , $(1.5, 0)$ , $L = 0.367$ (statement's printed answer: $L = 0.384$ )
plain Lloyd = scikit-learn Lloyd, no empty cluster	103 / 103
plain Lloyd = scikit-learn Lloyd, some cluster empties	6 / 197 (different conventions)
with k-means++ restarts, $m$ strictly above scikit-learn's	216 / 300
passes the rule $m \ge 0.95\, m_{\text{ref}}$ against scikit-learn	300 / 300

What to remember

Lloyd's algorithm is coordinate descent on $L$ : the assignment step is optimal for fixed centres, the mean is optimal for fixed clusters, so $L$ decreases until it converges.
It finds a local minimum only: initialisation matters, which is why practitioners use k-means++ and several restarts (n_init in scikit-learn).
Decide what happens to empty clusters: it is the classic source of disagreement between implementations.
Read the scoring rule before coding: here it rewards a lower loss, not an exact match, so restarts are free insurance.

Full code (solution.py)

"""FAIO 2024, problem D: K-means (Lloyd's algorithm from given initial centroids).

stdin : n, n points, k, k initial centroids (blank lines allowed)
stdout: the k final centroids, one "x y" per line

Scoring is by quality, not exact text: m = 1 - L / L_baseline (L = mean squared distance to the
nearest centroid, baseline = one centroid at the mean), and you pass if m >= 0.95 * m_reference.
We run Lloyd from the given centroids (what the reference does) and, as insurance, a few
k-means++ restarts; we output the solution with the lowest L, so m can only be >= the reference's.

Run:  python solution.py < input.txt
      python solution.py --check     (compare with scikit-learn on random tests)
      python solution.py --figure    (figures/lloyd.png)
"""
import random
import sys


def sqdist(p, c):
    return (p[0] - c[0]) ** 2 + (p[1] - c[1]) ** 2


def assign(points, cents):
    return [min(range(len(cents)), key=lambda j: sqdist(p, cents[j])) for p in points]


def loss(points, cents):
    return sum(min(sqdist(p, c) for c in cents) for p in points) / len(points)


def lloyd(points, cents, max_iter=1000, history=None):
    cents = [tuple(c) for c in cents]
    labels = None
    for _ in range(max_iter):
        new = assign(points, cents)                          # step 2: nearest centroid
        if new == labels:
            break                                            # converged: nothing moved
        labels = new
        if history is not None:
            history.append((list(cents), labels))
        sums = [[0.0, 0.0, 0] for _ in cents]
        for p, l in zip(points, labels):
            sums[l][0] += p[0]; sums[l][1] += p[1]; sums[l][2] += 1
        # step 3: move each centroid to the mean of its points (an empty cluster stays put)
        cents = [(sx / m, sy / m) if m else c for (sx, sy, m), c in zip(sums, cents)]
    return cents


def kmeanspp(points, k, rng):
    cents = [rng.choice(points)]
    while len(cents) < k:
        d = [min(sqdist(p, c) for c in cents) for p in points]
        cents.append(points[rng.choices(range(len(points)), weights=d)[0]] if sum(d) > 0
                     else rng.choice(points))
    return cents


def best_kmeans(points, init, restarts=5, seed=0):
    rng = random.Random(seed)
    best = lloyd(points, init)
    for _ in range(restarts):
        c = lloyd(points, kmeanspp(points, len(init), rng))
        if loss(points, c) < loss(points, best) - 1e-12:
            best = c
    return best


def parse(text):
    t = text.split()
    n = int(t[0]); pts = [(float(t[1 + 2 * i]), float(t[2 + 2 * i])) for i in range(n)]
    pos = 1 + 2 * n
    k = int(t[pos]); init = [(float(t[pos + 1 + 2 * j]), float(t[pos + 2 + 2 * j])) for j in range(k)]
    return pts, init


def fmt(v):
    s = f"{v:.6f}".rstrip("0").rstrip(".")
    return "0" if s in ("-0", "") else s


def solve():
    pts, init = parse(sys.stdin.read())
    for c in best_kmeans(pts, init):
        print(fmt(c[0]), fmt(c[1]))


def metric(points, cents):
    mx = sum(p[0] for p in points) / len(points); my = sum(p[1] for p in points) / len(points)
    base = loss(points, [(mx, my)])
    L = loss(points, cents)
    return 0.0 if L >= base else 1 - L / base


def check():
    import numpy as np
    from sklearn.cluster import KMeans
    pts, init = parse("5 -1 1 -1 2 -2 1 1 0 2 0 2 -2 2 3 0")
    print("sample:", [(fmt(x), fmt(y)) for x, y in lloyd(pts, init)],
          "| statement's output (1.5, 0), (-1.5, 1.3): L =", round(loss(pts, [(1.5, 0), (-1.5, 1.3)]), 4),
          "ours: L =", round(loss(pts, lloyd(pts, init)), 4))
    rng = random.Random(1)
    same = better = passed = 0
    T = 300
    for _ in range(T):
        k = rng.randint(2, 9); n = rng.randint(k, 999)
        centres = [(rng.uniform(5, 45), rng.uniform(5, 45)) for _ in range(rng.randint(1, 9))]
        pts = [(min(49.9, max(1.1, c[0] + rng.gauss(0, 3))), min(49.9, max(1.1, c[1] + rng.gauss(0, 3))))
               for c in (rng.choice(centres) for _ in range(n))]
        init = [(rng.uniform(1.1, 49.9), rng.uniform(1.1, 49.9)) for _ in range(k)]
        ours_plain = lloyd(pts, init)
        sk = KMeans(k, init=np.array(init), n_init=1, algorithm="lloyd", tol=0, max_iter=1000).fit(np.array(pts))
        m_ref = metric(pts, [tuple(c) for c in sk.cluster_centers_])
        same += abs(metric(pts, ours_plain) - m_ref) < 1e-9
        ours = best_kmeans(pts, init)
        better += metric(pts, ours) > m_ref + 1e-9
        passed += metric(pts, ours) >= 0.95 * m_ref - 1e-12
    print(f"{T} random tests: plain Lloyd == scikit-learn Lloyd on {same}; with restarts better on {better}; "
          f"pass (m >= 0.95 m_ref) on {passed}/{T}")


def figure():
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    rng = random.Random(7)
    centres = [(10, 10), (40, 10), (10, 40), (40, 40)]
    pts = [(c[0] + rng.gauss(0, 3), c[1] + rng.gauss(0, 3)) for c in (rng.choice(centres) for _ in range(400))]
    init = [(8, 10), (12, 10), (40, 12), (25, 40)]          # two centroids share one blob
    hist = []
    final = lloyd(pts, init, history=hist)
    cols = ["#1b9e77", "#d95f02", "#7570b3", "#e7298a"]
    fig, axes = plt.subplots(1, 3, figsize=(14, 4.6))
    for ax, it in zip(axes[:2], [0, len(hist) - 1]):
        cents, labels = hist[it]
        ax.scatter([p[0] for p in pts], [p[1] for p in pts], c=[cols[l] for l in labels], s=10, alpha=0.6)
        ax.scatter([c[0] for c in cents], [c[1] for c in cents], c=cols, marker="X", s=200, edgecolors="k")
        ax.set_title(f"iteration {it + 1}: assignment to the centroids"); ax.set_aspect("equal")
    for j in range(4):
        path = [h[0][j] for h in hist] + [final[j]]
        axes[1].plot([p[0] for p in path], [p[1] for p in path], color="k", lw=1, alpha=0.6)
    Ls = [loss(pts, h[0]) for h in hist] + [loss(pts, final)]
    axes[2].plot(range(1, len(Ls) + 1), Ls, "o-", color="#2171b5")
    best = best_kmeans(pts, init)
    axes[2].axhline(loss(pts, best), color="crimson", ls="--", label="best of k-means++ restarts")
    axes[2].set_xlabel("iteration"); axes[2].set_ylabel("L (mean squared distance)")
    axes[2].set_title("L never increases, yet can stall"); axes[2].legend()
    fig.tight_layout(); fig.savefig("figures/lloyd.png", dpi=150, facecolor="white")
    print("iterations", len(hist), "L final", loss(pts, final), "L best", loss(pts, best))


if __name__ == "__main__":
    if "--check" in sys.argv:
        check()
    elif "--figure" in sys.argv:
        figure()
    else:
        solve()