ckmeans.js

import makeMatrix from "./make_matrix.js";
import numericSort from "./numeric_sort.js";
import uniqueCountSorted from "./unique_count_sorted.js";

/**
 * 基于数据数组的累积和及平方累积和，生成增量计算值
 *
 * @private
 * @param {number} j 起始索引
 * @param {number} i 结束索引
 * @param {Array<number>} sums 累积和数组
 * @param {Array<number>} sumsOfSquares 平方累积和数组
 * @return {number} 计算得到的平方偏差值
 * @example
 * ssq(0, 1, [-1, 0, 2], [1, 1, 5]);
 */
function ssq(j, i, sums, sumsOfSquares) {
    let sji; // s(j, i) 表示从j到i的平方偏差和
    if (j > 0) {
        const muji = (sums[i] - sums[j - 1]) / (i - j + 1); // μ(j, i) 表示j到i的均值
        sji =
            sumsOfSquares[i] - sumsOfSquares[j - 1] - (i - j + 1) * muji * muji;
    } else {
        sji = sumsOfSquares[i] - (sums[i] * sums[i]) / (i + 1);
    }
    if (sji < 0) {
        return 0;
    }
    return sji;
}

/**
 * 递归执行分治策略计算指定簇的矩阵列
 *
 * @private
 * @param {number} iMin 当前簇计算的最小索引
 * @param {number} iMax 当前簇计算的最大索引
 * @param {number} cluster 当前计算的簇索引
 * @param {Array<Array<number>>} matrix 动态规划矩阵
 * @param {Array<Array<number>>} backtrackMatrix 回溯矩阵
 * @param {Array<number>} sums 累积和数组
 * @param {Array<number>} sumsOfSquares 平方累积和数组
 */
function fillMatrixColumn(
    iMin,
    iMax,
    cluster,
    matrix,
    backtrackMatrix,
    sums,
    sumsOfSquares
) {
    if (iMin > iMax) {
        return;
    }

    // 选取iMin和iMax之间的中点作为当前计算位置
    const i = Math.floor((iMin + iMax) / 2);

    matrix[cluster][i] = matrix[cluster - 1][i - 1];
    backtrackMatrix[cluster][i] = i;

    let jlow = cluster; // j的下界初始值
    if (iMin > cluster) {
        jlow = Math.max(jlow, backtrackMatrix[cluster][iMin - 1] || 0);
    }
    jlow = Math.max(jlow, backtrackMatrix[cluster - 1][i] || 0);

    let jhigh = i - 1; // j的上界初始值
    if (iMax < matrix[0].length - 1) {
        /* c8 ignore start */
        jhigh = Math.min(jhigh, backtrackMatrix[cluster][iMax + 1] || 0);
        /* c8 ignore end */
    }

    let sji;
    let sjlowi;
    let ssqjlow;
    let ssqj;
    for (let j = jhigh; j >= jlow; --j) {
        sji = ssq(j, i, sums, sumsOfSquares);

        if (sji + matrix[cluster - 1][jlow - 1] >= matrix[cluster][i]) {
            break;
        }

        // 计算簇边界的下界
        sjlowi = ssq(jlow, i, sums, sumsOfSquares);
        ssqjlow = sjlowi + matrix[cluster - 1][jlow - 1];

        if (ssqjlow < matrix[cluster][i]) {
            matrix[cluster][i] = ssqjlow;
            backtrackMatrix[cluster][i] = jlow;
        }
        jlow++;

        ssqj = sji + matrix[cluster - 1][j - 1];
        if (ssqj < matrix[cluster][i]) {
            matrix[cluster][i] = ssqj;
            backtrackMatrix[cluster][i] = j;
        }
    }

    // 递归处理左右子区间
    fillMatrixColumn(
        iMin,
        i - 1,
        cluster,
        matrix,
        backtrackMatrix,
        sums,
        sumsOfSquares
    );
    fillMatrixColumn(
        i + 1,
        iMax,
        cluster,
        matrix,
        backtrackMatrix,
        sums,
        sumsOfSquares
    );
}

/**
 * 初始化Ckmeans算法的主矩阵并启动分治聚类计算策略
 *
 * @private
 * @param {Array<number>} data 已排序的数值数组
 * @param {Array<Array<number>>} matrix 动态规划矩阵
 * @param {Array<Array<number>>} backtrackMatrix 回溯矩阵
 */
function fillMatrices(data, matrix, backtrackMatrix) {
    const nValues = matrix[0].length;

    // 通过中值平移提升数值稳定性
    const shift = data[Math.floor(nValues / 2)];

    // 计算平移后的累积和与平方累积和
    const sums = [];
    const sumsOfSquares = [];

    // 初始化矩阵和回溯矩阵的第一列（簇数为1时）
    for (let i = 0, shiftedValue; i < nValues; ++i) {
        shiftedValue = data[i] - shift;
        if (i === 0) {
            sums.push(shiftedValue);
            sumsOfSquares.push(shiftedValue * shiftedValue);
        } else {
            sums.push(sums[i - 1] + shiftedValue);
            sumsOfSquares.push(
                sumsOfSquares[i - 1] + shiftedValue * shiftedValue
            );
        }

        // 初始化簇数为0时的状态
        matrix[0][i] = ssq(0, i, sums, sumsOfSquares);
        backtrackMatrix[0][i] = 0;
    }

    // 初始化其他簇数的矩阵列
    let iMin;
    for (let cluster = 1; cluster < matrix.length; ++cluster) {
        if (cluster < matrix.length - 1) {
            iMin = cluster;
        } else {
            // 最后一列只需计算最后一个元素
            iMin = nValues - 1;
        }

        fillMatrixColumn(
            iMin,
            nValues - 1,
            cluster,
            matrix,
            backtrackMatrix,
            sums,
            sumsOfSquares
        );
    }
}

/**
 * Ckmeans聚类算法是一种基于动态规划的改进聚类方法，相比启发式的Jenks算法更为优化。
 * 该算法由[Haizhou Wang和Mingzhou Song](http://journal.r-project.org/archive/2011-2/RJournal_2011-2_Wang+Song.pdf)
 * 提出，通过最小化组内平方偏差和来实现数值数据的最优分组。
 *
 * 该实现基于Ckmeans 3.4.6版本，采用分治法将时间复杂度从O(kn²)优化至O(kn log(n))。
 *
 * ### 核心特性
 * - 动态规划构建代价矩阵和回溯矩阵
 * - 数值稳定性优化：通过中值平移技术
 * - 分治策略加速计算过程
 *
 * ### 参考文献
 * 《Ckmeans.1d.dp: Optimal k-means Clustering in One Dimension by Dynamic Programming》
 * 作者：Haizhou Wang和Mingzhou Song，刊于The R Journal Vol. 3/2, December 2011
 *
 * @param {Array<number>} x 输入数据集（数值数组）
 * @param {number} nClusters 期望的聚类数目（需小于等于数据长度）
 * @returns {Array<Array<number>>} 分组后的聚类结果
 * @throws {Error} 当请求聚类数大于数据长度时抛出异常
 * @example
 * ckmeans([-1, 2, -1, 2, 4, 5, 6, -1, 2, -1], 3);
 * // 返回：[[-1, -1, -1, -1], [2, 2, 2], [4, 5, 6]]
 */
function ckmeans(x, nClusters) {
    if (nClusters > x.length) {
        throw new Error("聚类数目不能超过数据元素个数");
    }

    const sorted = numericSort(x);
    const uniqueCount = uniqueCountSorted(sorted);

    // 处理全同数据情况
    if (uniqueCount === 1) {
        return [sorted];
    }

    const matrix = makeMatrix(nClusters, sorted.length);
    const backtrackMatrix = makeMatrix(nClusters, sorted.length);

    // 核心计算过程：填充动态规划矩阵和回溯矩阵
    fillMatrices(sorted, matrix, backtrackMatrix);

    // 通过回溯矩阵重建最优聚类分组
    const clusters = [];
    let clusterRight = backtrackMatrix[0].length - 1;

    // 从矩阵右下角开始逆向回溯
    for (let cluster = backtrackMatrix.length - 1; cluster >= 0; cluster--) {
        const clusterLeft = backtrackMatrix[cluster][clusterRight];
        clusters[cluster] = sorted.slice(clusterLeft, clusterRight + 1);

        if (cluster > 0) {
            clusterRight = clusterLeft - 1;
        }
    }

    return clusters;
}

export default ckmeans;