simple-squiggle

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

partitionSelect.js (4312B)

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