K-Means: Lloyd's Algorithm

// Lloyd's algorithm for k-means clustering on 2D points.
// Each frame: assign every point to nearest centroid, then move each
// centroid to the mean of its members. Converges when no assignment
// changes for ~30 consecutive frames.

let points = []; // {x, y, c}  -- c = current cluster index
let centroids = []; // {x, y}
let K = 4;
let stableFrames = 0;
let converged = false;
let cellBuf, cellCtx;
let cellImg, cellData, cellOwner;
const CELL_SCALE = 4; // voronoi raster downsampling
let cw = 0, ch = 0;
let lastBtnHit = -1; // for press-feedback on K +/- / R
const BTN_W = 44, BTN_H = 44;
// Scratch arrays for assignAndStep, sized to MAX_K to avoid per-frame allocs.
const MAX_K = 8;
const sxBuf = new Float64Array(MAX_K);
const syBuf = new Float64Array(MAX_K);
const ctBuf = new Int32Array(MAX_K);

function clusterColor(i, alpha = 1) {
  const h = (i * 360) / Math.max(K, 1);
  return `hsla(${h}, 75%, 60%, ${alpha})`;
}

function spawnBlobs(width, height) {
  points = [];
  const blobs = 4;
  const pad = 60;
  const centers = [];
  for (let b = 0; b < blobs; b++) {
    centers.push({
      x: pad + Math.random() * (width - 2 * pad),
      y: pad + Math.random() * (height - 2 * pad),
      sd: 18 + Math.random() * 22,
    });
  }
  const perBlob = Math.floor(300 / blobs);
  for (const c of centers) {
    for (let i = 0; i < perBlob; i++) {
      // Box-Muller for gaussian samples.
      const u1 = Math.max(1e-9, Math.random());
      const u2 = Math.random();
      const r = Math.sqrt(-2 * Math.log(u1));
      const t = 2 * Math.PI * u2;
      points.push({
        x: Math.max(2, Math.min(width - 2, c.x + r * Math.cos(t) * c.sd)),
        y: Math.max(2, Math.min(height - 2, c.y + r * Math.sin(t) * c.sd)),
        c: 0,
      });
    }
  }
}

function spawnCentroids(width, height) {
  centroids = [];
  for (let i = 0; i < K; i++) {
    centroids.push({
      x: 40 + Math.random() * (width - 80),
      y: 40 + Math.random() * (height - 80),
    });
  }
  stableFrames = 0;
  converged = false;
}

function ensureCellBuffers(width, height) {
  const nw = Math.max(1, Math.ceil(width / CELL_SCALE));
  const nh = Math.max(1, Math.ceil(height / CELL_SCALE));
  if (nw === cw && nh === ch && cellBuf) return;
  cw = nw; ch = nh;
  cellBuf = new OffscreenCanvas(cw, ch);
  cellCtx = cellBuf.getContext('2d');
  cellImg = cellCtx.createImageData(cw, ch);
  cellData = cellImg.data;
  cellOwner = new Int16Array(cw * ch);
}

function init({ canvas, ctx, width, height, input }) {
  K = 4;
  spawnBlobs(width, height);
  spawnCentroids(width, height);
  ensureCellBuffers(width, height);
}

function assignAndStep(width, height) {
  let changed = 0;
  // Assign step.
  for (const p of points) {
    let best = Infinity, bi = 0;
    for (let i = 0; i < K; i++) {
      const dx = p.x - centroids[i].x;
      const dy = p.y - centroids[i].y;
      const d = dx * dx + dy * dy;
      if (d < best) { best = d; bi = i; }
    }
    if (p.c !== bi) { p.c = bi; changed++; }
  }
  // Update step: mean of assigned points. Empty clusters reseed randomly.
  for (let i = 0; i < K; i++) { sxBuf[i] = 0; syBuf[i] = 0; ctBuf[i] = 0; }
  for (const p of points) {
    sxBuf[p.c] += p.x; syBuf[p.c] += p.y; ctBuf[p.c]++;
  }
  for (let i = 0; i < K; i++) {
    if (ctBuf[i] > 0) {
      // Ease toward target so motion is readable, not instant.
      const tx = sxBuf[i] / ctBuf[i];
      const ty = syBuf[i] / ctBuf[i];
      centroids[i].x += (tx - centroids[i].x) * 0.35;
      centroids[i].y += (ty - centroids[i].y) * 0.35;
    } else {
      centroids[i].x = 40 + Math.random() * (width - 80);
      centroids[i].y = 40 + Math.random() * (height - 80);
      changed++;
    }
  }
  if (changed === 0) stableFrames++; else stableFrames = 0;
  if (stableFrames >= 30) converged = true;
}

function drawVoronoi(ctx, width, height) {
  ensureCellBuffers(width, height);
  const data = cellData;
  const owner = cellOwner;
  for (let y = 0; y < ch; y++) {
    for (let x = 0; x < cw; x++) {
      const px = x * CELL_SCALE, py = y * CELL_SCALE;
      let best = Infinity, bi = 0;
      for (let i = 0; i < K; i++) {
        const dx = px - centroids[i].x, dy = py - centroids[i].y;
        const d = dx * dx + dy * dy;
        if (d < best) { best = d; bi = i; }
      }
      owner[y * cw + x] = bi;
    }
  }
  // Outline only (transparent fill) — draw edges where neighbor owner differs.
  for (let i = 3; i < data.length; i += 4) data[i] = 0;
  for (let y = 1; y < ch - 1; y++) {
    for (let x = 1; x < cw - 1; x++) {
      const i = y * cw + x;
      const o = owner[i];
      if (owner[i - 1] !== o || owner[i + 1] !== o || owner[i - cw] !== o || owner[i + cw] !== o) {
        const p = i * 4;
        data[p] = 255; data[p + 1] = 255; data[p + 2] = 255; data[p + 3] = 90;
      }
    }
  }
  cellCtx.putImageData(cellImg, 0, 0);
  ctx.imageSmoothingEnabled = false;
  ctx.drawImage(cellBuf, 0, 0, cw, ch, 0, 0, width, height);
}

function drawButton(ctx, x, y, label, pressed) {
  ctx.fillStyle = pressed ? 'rgba(255,255,255,0.22)' : 'rgba(255,255,255,0.10)';
  ctx.fillRect(x, y, BTN_W, BTN_H);
  ctx.strokeStyle = 'rgba(255,255,255,0.45)';
  ctx.lineWidth = 1;
  ctx.strokeRect(x + 0.5, y + 0.5, BTN_W - 1, BTN_H - 1);
  ctx.fillStyle = '#fff';
  ctx.font = 'bold 20px system-ui, sans-serif';
  ctx.textAlign = 'center'; ctx.textBaseline = 'middle';
  ctx.fillText(label, x + BTN_W / 2, y + BTN_H / 2);
}

function tick({ ctx, width, height, input }) {
  // Handle inputs.
  const clicks = input.consumeClicks();
  const minusX = 8, plusX = 8 + BTN_W + 6, resetX = 8 + (BTN_W + 6) * 2 + 16, btnY = height - BTN_H - 8;
  let pressedBtn = -1;
  for (const c of clicks) {
    if (c.y >= btnY && c.y <= btnY + BTN_H) {
      if (c.x >= minusX && c.x <= minusX + BTN_W) {
        if (K > 2) { K--; spawnCentroids(width, height); }
        pressedBtn = 0; continue;
      }
      if (c.x >= plusX && c.x <= plusX + BTN_W) {
        if (K < MAX_K) { K++; spawnCentroids(width, height); }
        pressedBtn = 1; continue;
      }
      if (c.x >= resetX && c.x <= resetX + BTN_W) {
        spawnCentroids(width, height);
        pressedBtn = 2; continue;
      }
    }
    // Else: add a new point at the click location.
    points.push({ x: c.x, y: c.y, c: 0 });
    stableFrames = 0; converged = false;
  }
  if (pressedBtn >= 0) lastBtnHit = pressedBtn;
  const jp = input.justPressed;
  if (jp && (jp('r') || jp('R'))) spawnCentroids(width, height);
  if (jp && (jp('k') || jp('K'))) {
    K = K >= MAX_K ? 2 : K + 1;
    spawnCentroids(width, height);
  }

  assignAndStep(width, height);

  // Render.
  ctx.fillStyle = '#0a0a12';
  ctx.fillRect(0, 0, width, height);
  drawVoronoi(ctx, width, height);

  // Points.
  for (const p of points) {
    ctx.fillStyle = clusterColor(p.c, 0.85);
    ctx.beginPath(); ctx.arc(p.x, p.y, 2.2, 0, Math.PI * 2); ctx.fill();
  }
  // Centroids.
  for (let i = 0; i < K; i++) {
    ctx.fillStyle = clusterColor(i, 1);
    ctx.strokeStyle = '#fff'; ctx.lineWidth = 2;
    ctx.beginPath(); ctx.arc(centroids[i].x, centroids[i].y, 9, 0, Math.PI * 2);
    ctx.fill(); ctx.stroke();
  }

  // HUD.
  ctx.fillStyle = 'rgba(255,255,255,0.85)';
  ctx.font = '12px system-ui, sans-serif';
  ctx.textAlign = 'left'; ctx.textBaseline = 'top';
  ctx.fillText(`k = ${K}   n = ${points.length}   stable = ${stableFrames}`, 10, 10);
  if (converged) {
    ctx.fillStyle = 'rgba(80, 220, 120, 0.95)';
    ctx.fillRect(width - 110, 8, 100, 22);
    ctx.fillStyle = '#06120a';
    ctx.font = 'bold 12px system-ui, sans-serif';
    ctx.textAlign = 'center'; ctx.textBaseline = 'middle';
    ctx.fillText('converged', width - 60, 19);
  }

  // K buttons + reset.
  drawButton(ctx, minusX, btnY, '−', lastBtnHit === 0);
  drawButton(ctx, plusX, btnY, '+', lastBtnHit === 1);
  drawButton(ctx, resetX, btnY, 'R', lastBtnHit === 2);
  ctx.textAlign = 'left'; ctx.textBaseline = 'top';
  ctx.fillStyle = 'rgba(255,255,255,0.65)';
  ctx.font = '11px system-ui, sans-serif';
  ctx.fillText('k', plusX + BTN_W + 4, btnY + BTN_H / 2 - 6);
}