kd.js

/**
 * kd-tree based clustering
 *
 * Based on kdbush package
 */

'use strict'


module.exports = function clusterKD (points, ids, levels, weights, options) {
    let ptr = 0
    let n = ids.length
    let nodeSize = options.nodeSize

    sort(0, n - 1, 0)

    function sort(left, right, level) {
        let count = right - left
        weights[ptr] = count
        levels[ptr++] = level

        if (count <= nodeSize) return;

        let m = Math.floor((left + right) / 2);

        select(m, left, right, level % 2);

        sort(left, m - 1, level + 1);
        sort(m + 1, right, level + 1);
    }

    function select(k, left, right, inc) {
        while (right > left) {
            if (right - left > 600) {
                let n = right - left + 1;
                let m = k - left + 1;
                let z = Math.log(n);
                let s = 0.5 * Math.exp(2 * z / 3);
                let sd = 0.5 * Math.sqrt(z * s * (n - s) / n) * (m - n / 2 < 0 ? -1 : 1);
                let newLeft = Math.max(left, Math.floor(k - m * s / n + sd));
                let newRight = Math.min(right, Math.floor(k + (n - m) * s / n + sd));
                select(k, newLeft, newRight, inc);
            }

            let t = points[2 * k + inc];
            let i = left;
            let j = right;

            swapItem(left, k);
            if (points[2 * right + inc] > t) swapItem(left, right);

            while (i < j) {
                swapItem(i, j);
                i++;
                j--;
                while (points[2 * i + inc] < t) i++;
                while (points[2 * j + inc] > t) j--;
            }

            if (points[2 * left + inc] === t) swapItem(left, j);
            else {
                j++;
                swapItem(j, right);
            }

            if (j <= k) left = j + 1;
            if (k <= j) right = j - 1;
        }
    }

    function swapItem(i, j) {
        swap(ids, i, j);
        swap(points, 2 * i, 2 * j);
        swap(points, 2 * i + 1, 2 * j + 1);
    }

    function swap(arr, i, j) {
        let tmp = arr[i];
        arr[i] = arr[j];
        arr[j] = tmp;
    }
}