k meansをbigqueryでやろう

CREATE TEMP FUNCTION kmeans1D(values ARRAY<FLOAT64>, k INT64, max_iter INT64)
RETURNS ARRAY<STRUCT<value INT64, group_id INT64>>
LANGUAGE js AS """
  function getInitialCenters(vals, k) {
    let sorted = vals.slice().sort((a, b) => a - b);  // 非破壊ソート
    const step = Math.floor(sorted.length / (k + 1));
    let centers = [];
    for (let i = 1; i <= k; i++) {
      const idx = Math.min(i * step, sorted.length - 1);
      centers.push(sorted[idx]);
    }
    return centers;
  }

  function assignGroups(vals, centers) {
    return vals.map(val => {
      const distances = centers.map(c => Math.pow(val - c, 2));
      return distances.indexOf(Math.min(...distances));
    });
  }

  function updateCenters(vals, labels, k, prevCenters) {
    let sums = new Array(k).fill(0);
    let counts = new Array(k).fill(0);
    for (let i = 0; i < vals.length; i++) {
      const label = labels[i];
      sums[label] += vals[i];
      counts[label]++;
    }
    return sums.map((sum, i) =>
      counts[i] > 0 ? sum / counts[i] : prevCenters[i]
    );
  }

  let centers = getInitialCenters(values, k);
  let labels = [];

  for (let iter = 0; iter < max_iter; iter++) {
    labels = assignGroups(values, centers);
    const newCenters = updateCenters(values, labels, k, centers);
    const changed = centers.some((c, i) => Math.abs(c - newCenters[i]) > 1e-6);
    if (!changed) break;
    centers = newCenters;
  }

  return values.map((val, idx) => {
    return { value: Math.round(val), group_id: labels[idx] };
  });
""";