simple-squiggle

A restricted subset of Squiggle
Log | Files | Refs | README

partitionSelect.js (4030B)

      1 import { isMatrix } from '../../utils/is.js';
      2 import { isInteger } from '../../utils/number.js';
      3 import { factory } from '../../utils/factory.js';
      4 var name = 'partitionSelect';
      5 var dependencies = ['typed', 'isNumeric', 'isNaN', 'compare'];
      6 export var createPartitionSelect = /* #__PURE__ */factory(name, dependencies, _ref => {
      7   var {
      8     typed,
      9     isNumeric,
     10     isNaN,
     11     compare
     12   } = _ref;
     13   var asc = compare;
     14 
     15   var desc = (a, b) => -compare(a, b);
     16   /**
     17    * Partition-based selection of an array or 1D matrix.
     18    * Will find the kth smallest value, and mutates the input array.
     19    * Uses Quickselect.
     20    *
     21    * Syntax:
     22    *
     23    *    math.partitionSelect(x, k)
     24    *    math.partitionSelect(x, k, compare)
     25    *
     26    * Examples:
     27    *
     28    *    math.partitionSelect([5, 10, 1], 2)           // returns 10
     29    *    math.partitionSelect(['C', 'B', 'A', 'D'], 1) // returns 'B'
     30    *
     31    *    function sortByLength (a, b) {
     32    *      return a.length - b.length
     33    *    }
     34    *    math.partitionSelect(['Langdon', 'Tom', 'Sara'], 2, sortByLength) // returns 'Langdon'
     35    *
     36    * See also:
     37    *
     38    *    sort
     39    *
     40    * @param {Matrix | Array} x    A one dimensional matrix or array to sort
     41    * @param {Number} k            The kth smallest value to be retrieved zero-based index
     42    * @param {Function | 'asc' | 'desc'} [compare='asc']
     43    *        An optional comparator function. The function is called as
     44    *        `compare(a, b)`, and must return 1 when a > b, -1 when a < b,
     45    *        and 0 when a == b.
     46    * @return {*} Returns the kth lowest value.
     47    */
     48 
     49 
     50   return typed(name, {
     51     'Array | Matrix, number': function ArrayMatrixNumber(x, k) {
     52       return _partitionSelect(x, k, asc);
     53     },
     54     'Array | Matrix, number, string': function ArrayMatrixNumberString(x, k, compare) {
     55       if (compare === 'asc') {
     56         return _partitionSelect(x, k, asc);
     57       } else if (compare === 'desc') {
     58         return _partitionSelect(x, k, desc);
     59       } else {
     60         throw new Error('Compare string must be "asc" or "desc"');
     61       }
     62     },
     63     'Array | Matrix, number, function': _partitionSelect
     64   });
     65 
     66   function _partitionSelect(x, k, compare) {
     67     if (!isInteger(k) || k < 0) {
     68       throw new Error('k must be a non-negative integer');
     69     }
     70 
     71     if (isMatrix(x)) {
     72       var size = x.size();
     73 
     74       if (size.length > 1) {
     75         throw new Error('Only one dimensional matrices supported');
     76       }
     77 
     78       return quickSelect(x.valueOf(), k, compare);
     79     }
     80 
     81     if (Array.isArray(x)) {
     82       return quickSelect(x, k, compare);
     83     }
     84   }
     85   /**
     86    * Quickselect algorithm.
     87    * Code adapted from:
     88    * https://blog.teamleadnet.com/2012/07/quick-select-algorithm-find-kth-element.html
     89    *
     90    * @param {Array} arr
     91    * @param {Number} k
     92    * @param {Function} compare
     93    * @private
     94    */
     95 
     96 
     97   function quickSelect(arr, k, compare) {
     98     if (k >= arr.length) {
     99       throw new Error('k out of bounds');
    100     } // check for NaN values since these can cause an infinite while loop
    101 
    102 
    103     for (var i = 0; i < arr.length; i++) {
    104       if (isNumeric(arr[i]) && isNaN(arr[i])) {
    105         return arr[i]; // return NaN
    106       }
    107     }
    108 
    109     var from = 0;
    110     var to = arr.length - 1; // if from == to we reached the kth element
    111 
    112     while (from < to) {
    113       var r = from;
    114       var w = to;
    115       var pivot = arr[Math.floor(Math.random() * (to - from + 1)) + from]; // stop if the reader and writer meets
    116 
    117       while (r < w) {
    118         // arr[r] >= pivot
    119         if (compare(arr[r], pivot) >= 0) {
    120           // put the large values at the end
    121           var tmp = arr[w];
    122           arr[w] = arr[r];
    123           arr[r] = tmp;
    124           --w;
    125         } else {
    126           // the value is smaller than the pivot, skip
    127           ++r;
    128         }
    129       } // if we stepped up (r++) we need to step one down (arr[r] > pivot)
    130 
    131 
    132       if (compare(arr[r], pivot) > 0) {
    133         --r;
    134       } // the r pointer is on the end of the first k elements
    135 
    136 
    137       if (k <= r) {
    138         to = r;
    139       } else {
    140         from = r + 1;
    141       }
    142     }
    143 
    144     return arr[k];
    145   }
    146 });