Decision Tree: Greedy Gini Splits

// Greedy decision-tree classifier on 2D points (2 classes).
// Each internal node picks the axis-aligned split minimizing weighted child
// Gini. Tree grows depth-by-depth: left panel shows feature space with cut
// lines and class-shaded regions; right panel renders the tree. L-click adds
// red (class 0), R-click adds blue (class 1); the tree rebuilds.

const MAX_D = 5, MIN_S = 3, GROW = 24, RS = 4;
let pts = [], root = null, rd = 0, fsg = 0;
let W = 0, H = 0, rbuf = null, rctx = null;
// Current class for taps (mobile / single-button input). Toggle via on-canvas
// button. Desktop L/R-click still picks 0/1 directly regardless of this.
let curClass = 0;
// On-canvas button rects, updated each frame to track resize.
let btnToggle = { x: 0, y: 0, w: 0, h: 0 };
let btnReset = { x: 0, y: 0, w: 0, h: 0 };
function inRect(r, x, y) { return x >= r.x && x <= r.x + r.w && y >= r.y && y <= r.y + r.h; }

function gini(a, b) { const n = a + b; if (!n) return 0; const p = a / n, q = b / n; return 1 - p * p - q * q; }
function maj(a) { let c0 = 0, c1 = 0; for (const p of a) (p.c ? c1++ : c0++); return c0 >= c1 ? 0 : 1; }
function counts(a) { let c0 = 0, c1 = 0; for (const p of a) (p.c ? c1++ : c0++); return [c0, c1]; }

function bestSplit(a) {
  if (a.length < 2 * MIN_S) return null;
  const [c0, c1] = counts(a), N = a.length, baseG = gini(c0, c1);
  if (baseG === 0) return null;
  let best = null;
  for (let ax = 0; ax < 2; ax++) {
    const s = a.slice().sort((p, q) => (ax ? p.y - q.y : p.x - q.x));
    let lc0 = 0, lc1 = 0;
    for (let i = 0; i < N - 1; i++) {
      const p = s[i]; (p.c ? lc1++ : lc0++);
      const v = ax ? p.y : p.x, vn = ax ? s[i + 1].y : s[i + 1].x;
      if (v === vn) continue;
      const lN = i + 1, rN = N - lN;
      if (lN < MIN_S || rN < MIN_S) continue;
      const g = (lN * gini(lc0, lc1) + rN * gini(c0 - lc0, c1 - lc1)) / N;
      const gain = baseG - g;
      if (gain > 1e-9 && (!best || gain > best.gain)) best = { ax, thr: (v + vn) / 2, gain };
    }
  }
  if (!best) return null;
  const L = [], R = [];
  for (const p of a) ((best.ax ? p.y : p.x) <= best.thr ? L : R).push(p);
  return { ax: best.ax, thr: best.thr, L, R };
}

function build(arr, depth, bb) {
  const node = { depth, pts: arr, n: arr.length, label: maj(arr), bbox: bb, leaf: true, ax: -1, thr: 0, L: null, R: null };
  if (!arr.length || depth >= MAX_D) return node;
  const [c0, c1] = counts(arr);
  if (gini(c0, c1) === 0) return node;
  const sp = bestSplit(arr);
  if (!sp) return node;
  node.leaf = false; node.ax = sp.ax; node.thr = sp.thr;
  const lb = sp.ax === 0 ? { x0: bb.x0, y0: bb.y0, x1: sp.thr, y1: bb.y1 } : { x0: bb.x0, y0: bb.y0, x1: bb.x1, y1: sp.thr };
  const rb = sp.ax === 0 ? { x0: sp.thr, y0: bb.y0, x1: bb.x1, y1: bb.y1 } : { x0: bb.x0, y0: sp.thr, x1: bb.x1, y1: bb.y1 };
  node.L = build(sp.L, depth + 1, lb);
  node.R = build(sp.R, depth + 1, rb);
  return node;
}

const lpw = () => Math.floor(W * 0.62);

function rebuild() {
  root = build(pts.slice(), 0, { x0: 0, y0: 0, x1: lpw(), y1: H });
  rd = 0; fsg = 0;
}

function seed() {
  pts = [];
  const w = lpw(), s = Math.min(w, H) * 0.07;
  function blob(cx, cy, sd, cls, n) {
    for (let i = 0; i < n; i++) {
      const u1 = Math.max(1e-9, Math.random()), u2 = Math.random();
      const r = Math.sqrt(-2 * Math.log(u1)) * sd, t = 2 * Math.PI * u2;
      pts.push({
        x: Math.max(4, Math.min(w - 4, cx + r * Math.cos(t))),
        y: Math.max(4, Math.min(H - 4, cy + r * Math.sin(t))),
        c: cls,
      });
    }
  }
  blob(w * 0.28, H * 0.30, s, 0, 18);
  blob(w * 0.72, H * 0.30, s, 0, 16);
  blob(w * 0.50, H * 0.72, s * 1.15, 1, 22);
  rebuild();
}

function init({ width, height }) {
  W = width; H = height;
  seed();
  const rw = Math.max(1, Math.ceil(lpw() / RS)), rh = Math.max(1, Math.ceil(H / RS));
  rbuf = new OffscreenCanvas(rw, rh); rctx = rbuf.getContext('2d');
}

function leafFor(n, px, py, md) {
  while (n && !n.leaf && n.depth < md) n = (n.ax ? py : px) <= n.thr ? n.L : n.R;
  return n;
}

function drawRegions(ctx) {
  const w = lpw();
  const rw = Math.max(1, Math.ceil(w / RS)), rh = Math.max(1, Math.ceil(H / RS));
  if (rbuf.width !== rw || rbuf.height !== rh) { rbuf = new OffscreenCanvas(rw, rh); rctx = rbuf.getContext('2d'); }
  const img = rctx.createImageData(rw, rh), d = img.data;
  for (let yy = 0; yy < rh; yy++) for (let xx = 0; xx < rw; xx++) {
    const lf = leafFor(root, xx * RS, yy * RS, rd), i = (yy * rw + xx) * 4;
    if (!lf || lf.n === 0) { d[i + 3] = 0; continue; }
    const [c0, c1] = counts(lf.pts), p1 = c1 / (c0 + c1), pur = Math.abs(p1 - 0.5) * 2;
    d[i] = Math.round(255 * (1 - p1) * 0.85 + 30 * p1);
    d[i + 1] = 40;
    d[i + 2] = Math.round(255 * p1 * 0.85 + 30 * (1 - p1));
    d[i + 3] = Math.round(40 + pur * 90);
  }
  rctx.putImageData(img, 0, 0);
  ctx.imageSmoothingEnabled = false;
  ctx.drawImage(rbuf, 0, 0, rw, rh, 0, 0, w, H);
}

function drawCuts(ctx) {
  (function walk(n) {
    if (!n || n.leaf || n.depth >= rd) return;
    ctx.strokeStyle = `rgba(255,255,255,${0.35 + 0.65 * (1 - n.depth / MAX_D)})`;
    ctx.lineWidth = n.depth === 0 ? 2 : 1.25;
    ctx.setLineDash(n.depth === 0 ? [] : [4, 3]);
    ctx.beginPath();
    if (n.ax === 0) { ctx.moveTo(n.thr, n.bbox.y0); ctx.lineTo(n.thr, n.bbox.y1); }
    else { ctx.moveTo(n.bbox.x0, n.thr); ctx.lineTo(n.bbox.x1, n.thr); }
    ctx.stroke(); ctx.setLineDash([]);
    walk(n.L); walk(n.R);
  })(root);
}

function drawPoints(ctx) {
  for (const p of pts) {
    ctx.fillStyle = p.c ? "#5aa8ff" : "#ff5a6a";
    ctx.beginPath(); ctx.arc(p.x, p.y, 3.2, 0, Math.PI * 2); ctx.fill();
    ctx.strokeStyle = "rgba(0,0,0,0.55)"; ctx.lineWidth = 0.8; ctx.stroke();
  }
}

function layout(node, ox, pw, ph) {
  const leaves = [];
  (function collect(n) { if (!n) return; if (n.leaf || n.depth >= rd) { leaves.push(n); return; } collect(n.L); collect(n.R); })(node);
  const vmd = Math.min(rd, MAX_D);
  const dy = Math.max(28, (ph - 40) / Math.max(1, vmd + 1));
  const dx = (pw - 30) / Math.max(1, leaves.length);
  const pos = new Map();
  leaves.forEach((lf, i) => pos.set(lf, ox + 15 + (i + 0.5) * dx));
  (function assign(n) { if (!n) return 0; if (pos.has(n)) return pos.get(n); const x = (assign(n.L) + assign(n.R)) / 2; pos.set(n, x); return x; })(node);
  return { pos, dy };
}

function drawTree(ctx) {
  const ox = lpw(), pw = W - ox;
  ctx.fillStyle = "#0b0d18"; ctx.fillRect(ox, 0, pw, H);
  ctx.fillStyle = "rgba(200,210,235,0.85)";
  ctx.font = "12px ui-sans-serif, system-ui, sans-serif";
  ctx.fillText(`Depth ${rd}/${MAX_D}`, ox + 10, 16);
  if (!root || root.n === 0) return;
  const { pos, dy } = layout(root, ox, pw, H);
  const yOf = (d) => 30 + d * dy;
  (function edges(n) {
    if (!n || n.leaf || n.depth >= rd) return;
    const px = pos.get(n), py = yOf(n.depth);
    for (const ch of [n.L, n.R]) {
      ctx.strokeStyle = "rgba(180,190,220,0.55)"; ctx.lineWidth = 1;
      ctx.beginPath(); ctx.moveTo(px, py); ctx.lineTo(pos.get(ch), yOf(ch.depth)); ctx.stroke();
    }
    edges(n.L); edges(n.R);
  })(root);
  (function nodes(n) {
    if (!n) return;
    const x = pos.get(n), y = yOf(n.depth), vl = n.leaf || n.depth >= rd;
    const [c0, c1] = counts(n.pts), p1 = n.n ? c1 / n.n : 0.5;
    ctx.fillStyle = `rgb(${Math.round(255 * (1 - p1) * 0.85 + 30 * p1)},40,${Math.round(255 * p1 * 0.85 + 30 * (1 - p1))})`;
    const r = vl ? 9 : 6;
    ctx.beginPath(); ctx.arc(x, y, r, 0, Math.PI * 2); ctx.fill();
    ctx.strokeStyle = vl ? "#fff" : "rgba(255,255,255,0.55)"; ctx.lineWidth = vl ? 1.5 : 1; ctx.stroke();
    ctx.fillStyle = "rgba(220,230,255,0.85)";
    ctx.font = "10px ui-monospace, monospace"; ctx.textAlign = "center";
    if (!vl) ctx.fillText((n.ax ? "y" : "x") + " ≤ " + n.thr.toFixed(0), x, y - r - 4);
    else ctx.fillText(`${c0}/${c1}`, x, y + r + 11);
    ctx.textAlign = "left";
    if (!n.leaf && n.depth < rd) { nodes(n.L); nodes(n.R); }
  })(root);
}

function leafCount(n, md) { if (!n) return 0; if (n.leaf || n.depth >= md) return 1; return leafCount(n.L, md) + leafCount(n.R, md); }
function acc() { if (!root || !pts.length) return 1; let k = 0; for (const p of pts) { const lf = leafFor(root, p.x, p.y, rd); if (lf && lf.label === p.c) k++; } return k / pts.length; }

function tick({ ctx, width, height, input }) {
  if (width !== W || height !== H) { W = width; H = height; rebuild(); }
  const w = lpw();
  // Refresh button hitboxes before consuming clicks so first-frame taps work.
  const bbw = 78, bbh = 26, bbx = w - 8 - bbw, bby = 8;
  btnToggle = { x: bbx, y: bby, w: bbw, h: bbh };
  btnReset  = { x: bbx, y: bby + bbh + 6, w: bbw, h: bbh };
  let dirty = false, didReset = false;
  for (const c of input.consumeClicks()) {
    if (inRect(btnToggle, c.x, c.y)) { curClass = curClass ? 0 : 1; continue; }
    if (inRect(btnReset, c.x, c.y)) { seed(); didReset = true; continue; }
    if (c.x < 0 || c.x > w || c.y < 0 || c.y > H) continue;
    if (pts.length >= 200) pts.shift();
    const cls = c.button === 2 ? 1 : (c.button === 0 ? curClass : 0);
    pts.push({ x: c.x, y: c.y, c: cls });
    dirty = true;
  }
  if (input.justPressed && (input.justPressed("r") || input.justPressed("R"))) { seed(); didReset = true; }
  if (dirty && !didReset) rebuild();

  fsg++;
  if (rd < MAX_D && fsg >= GROW) {
    const more = (function deeper(n) { if (!n || n.leaf) return false; if (n.depth >= rd) return true; return deeper(n.L) || deeper(n.R); })(root);
    if (more) rd++;
    fsg = 0;
  }

  ctx.fillStyle = "#06070d"; ctx.fillRect(0, 0, w, H);
  drawRegions(ctx); drawCuts(ctx); drawPoints(ctx);
  ctx.strokeStyle = "#1a1d2a"; ctx.lineWidth = 2;
  ctx.beginPath(); ctx.moveTo(w, 0); ctx.lineTo(w, H); ctx.stroke();
  drawTree(ctx);

  ctx.fillStyle = "rgba(0,0,0,0.55)"; ctx.fillRect(8, 8, 178, 56);
  ctx.fillStyle = "#fff"; ctx.font = "12px ui-monospace, monospace";
  ctx.textAlign = "left"; ctx.textBaseline = "alphabetic";
  ctx.fillText(`n = ${pts.length}`, 16, 24);
  ctx.fillText(`leaves = ${leafCount(root, rd)}`, 16, 40);
  ctx.fillText(`train acc = ${(acc() * 100).toFixed(1)}%`, 16, 56);

  // On-canvas buttons (top-right of the scatter panel). Tracks current class
  // for taps and offers a reset that works on mobile.
  ctx.fillStyle = curClass ? "rgba(90,168,255,0.85)" : "rgba(255,90,106,0.85)";
  ctx.fillRect(btnToggle.x, btnToggle.y, btnToggle.w, btnToggle.h);
  ctx.fillStyle = "rgba(255,255,255,0.92)";
  ctx.font = "12px ui-sans-serif, system-ui, sans-serif";
  ctx.textAlign = "center"; ctx.textBaseline = "middle";
  ctx.fillText(curClass ? "Tap: blue" : "Tap: red", btnToggle.x + btnToggle.w / 2, btnToggle.y + btnToggle.h / 2);
  ctx.fillStyle = "rgba(40,46,66,0.85)";
  ctx.fillRect(btnReset.x, btnReset.y, btnReset.w, btnReset.h);
  ctx.fillStyle = "rgba(220,230,255,0.95)";
  ctx.fillText("Reset", btnReset.x + btnReset.w / 2, btnReset.y + btnReset.h / 2);
  ctx.textAlign = "left"; ctx.textBaseline = "alphabetic";

  ctx.fillStyle = "rgba(200,220,255,0.7)";
  ctx.font = "11px system-ui, sans-serif";
  ctx.fillText("tap to add · toggle class · L/R-click on desktop", 10, H - 10);
}