partitionSelect.js 4.2 KB

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