partitionSelect.js 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151
  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, math.compareText) // 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. * // the input array is mutated
  41. * arr = [5, 2, 1]
  42. * math.partitionSelect(arr, 0) // returns 1, arr is now: [1, 2, 5]
  43. * math.partitionSelect(arr, 1, 'desc') // returns 2, arr is now: [5, 2, 1]
  44. *
  45. * See also:
  46. *
  47. * sort
  48. *
  49. * @param {Matrix | Array} x A one dimensional matrix or array to sort
  50. * @param {Number} k The kth smallest value to be retrieved zero-based index
  51. * @param {Function | 'asc' | 'desc'} [compare='asc']
  52. * An optional comparator function. The function is called as
  53. * `compare(a, b)`, and must return 1 when a > b, -1 when a < b,
  54. * and 0 when a == b.
  55. * @return {*} Returns the kth lowest value.
  56. */
  57. return typed(name, {
  58. 'Array | Matrix, number': function ArrayMatrixNumber(x, k) {
  59. return _partitionSelect(x, k, asc);
  60. },
  61. 'Array | Matrix, number, string': function ArrayMatrixNumberString(x, k, compare) {
  62. if (compare === 'asc') {
  63. return _partitionSelect(x, k, asc);
  64. } else if (compare === 'desc') {
  65. return _partitionSelect(x, k, desc);
  66. } else {
  67. throw new Error('Compare string must be "asc" or "desc"');
  68. }
  69. },
  70. 'Array | Matrix, number, function': _partitionSelect
  71. });
  72. function _partitionSelect(x, k, compare) {
  73. if (!(0, _number.isInteger)(k) || k < 0) {
  74. throw new Error('k must be a non-negative integer');
  75. }
  76. if ((0, _is.isMatrix)(x)) {
  77. var size = x.size();
  78. if (size.length > 1) {
  79. throw new Error('Only one dimensional matrices supported');
  80. }
  81. return quickSelect(x.valueOf(), k, compare);
  82. }
  83. if (Array.isArray(x)) {
  84. return quickSelect(x, k, compare);
  85. }
  86. }
  87. /**
  88. * Quickselect algorithm.
  89. * Code adapted from:
  90. * https://blog.teamleadnet.com/2012/07/quick-select-algorithm-find-kth-element.html
  91. *
  92. * @param {Array} arr
  93. * @param {Number} k
  94. * @param {Function} compare
  95. * @private
  96. */
  97. function quickSelect(arr, k, compare) {
  98. if (k >= arr.length) {
  99. throw new Error('k out of bounds');
  100. }
  101. // check for NaN values since these can cause an infinite while loop
  102. for (var i = 0; i < arr.length; i++) {
  103. if (isNumeric(arr[i]) && isNaN(arr[i])) {
  104. return arr[i]; // return NaN
  105. }
  106. }
  107. var from = 0;
  108. var to = arr.length - 1;
  109. // if from == to we reached the kth element
  110. while (from < to) {
  111. var r = from;
  112. var w = to;
  113. var pivot = arr[Math.floor(Math.random() * (to - from + 1)) + from];
  114. // stop if the reader and writer meets
  115. while (r < w) {
  116. // arr[r] >= pivot
  117. if (compare(arr[r], pivot) >= 0) {
  118. // put the large values at the end
  119. var tmp = arr[w];
  120. arr[w] = arr[r];
  121. arr[r] = tmp;
  122. --w;
  123. } else {
  124. // the value is smaller than the pivot, skip
  125. ++r;
  126. }
  127. }
  128. // if we stepped up (r++) we need to step one down (arr[r] > pivot)
  129. if (compare(arr[r], pivot) > 0) {
  130. --r;
  131. }
  132. // the r pointer is on the end of the first k elements
  133. if (k <= r) {
  134. to = r;
  135. } else {
  136. from = r + 1;
  137. }
  138. }
  139. return arr[k];
  140. }
  141. });
  142. exports.createPartitionSelect = createPartitionSelect;