partitionSelect.js 3.9 KB

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