diff --git a/lib/function/utils/partitionSelect.js b/lib/function/utils/partitionSelect.js new file mode 100644 index 000000000..b3c11939d --- /dev/null +++ b/lib/function/utils/partitionSelect.js @@ -0,0 +1,143 @@ +'use strict'; + +module.exports = function (math) { + var util = require('../../util/index'); + var Matrix = math.type.Matrix; + + var isNumber = util.number.isNumber; + var isInteger = util.number.isInteger; + + /** + * Partition-based selection of an array or 1D matrix. + * Will find the kth smallest value, and mutates the input array. + * Currently uses Quickselect. + * + * Syntax: + * + * math.partitionSelect(x, k) + * math.partitionSelect(x, k, compare) + * + * Examples: + * + * math.partitionSelect([5, 10, 1], 2); // returns 10 + * math.partitionSelect(['C', 'B', 'A', 'D'], 1); // returns 'B' + * + * function sortByLength (a, b) { + * return a.length - b.length; + * } + * math.partitionSelect(['Langdon', 'Tom', 'Sara'], 2, sortByLength); // returns 'Langdon' + * + * See also: + * + * sort + * + * @param {Matrix | Array} x A one dimensional matrix or array to sort + * @param {Number} k The kth smallest value to be retrieved; zero-based index + * @param {Function | 'asc' | 'desc'} [compare='asc'] + * An optional comparator function. The function is called as + * `compare(a, b)`, and must return 1 when a > b, -1 when a < b, + * and 0 when a == b. + * @return {*} Returns the kth lowest value. + */ + math.partitionSelect = function (x, k, compare) { + var _compare; + + if (arguments.length === 2) { + _compare = math.compare; + } + else if (arguments.length === 3) { + if (typeof compare === 'function') { + _compare = compare; + } + else if (compare === 'asc') { + _compare = math.compare; + } + else if (compare === 'desc') { + _compare = function (a, b) { + return -math.compare(a, b); + } + } + else { + throw new math.error.UnsupportedTypeError('partitionSelect', math['typeof'](compare)); + } + } + else { + throw new math.error.ArgumentsError('partitionSelect', arguments.length, 2, 3); + } + + if (isNumber(k)) { + if (isInteger(k) && k >= 0) { + if (x instanceof Matrix) { + var size = x.size(); + if (size.length > 1) { + throw new Error('Only one dimensional matrices supported'); + } + return quickselect(x.valueOf(), k, _compare); + } + + if (Array.isArray(x)) { + return quickselect(x, k, _compare); + } + + throw new math.error.UnsupportedTypeError('partitionSelect', math['typeof'](x)); + } + + throw new Error('k must be a non-negative integer'); + } + + throw new math.error.UnsupportedTypeError('partitionSelect', math['typeof'](k)); + }; + + /** + * Quickselect algorithm. + * Code adapted from: + * http://blog.teamleadnet.com/2012/07/quick-select-algorithm-find-kth-element.html + * + * @param {Array} arr + * @param {Number} k + * @param {Function} compare + * @private + */ + function quickselect(arr, k, compare) { + if (k >= arr.length) { + throw new Error('k out of bounds'); + } + + var from = 0; + var to = arr.length - 1; + + // if from == to we reached the kth element + while (from < to) { + var r = from; + var w = to; + var pivot = arr[Math.floor(Math.random() * (to - from + 1)) + from]; + + // stop if the reader and writer meets + while (r < w) { + // arr[r] >= pivot + if (compare(arr[r], pivot) >= 0) { // put the large values at the end + var tmp = arr[w]; + arr[w] = arr[r]; + arr[r] = tmp; + --w; + } else { // the value is smaller than the pivot, skip + ++r; + } + } + + // if we stepped up (r++) we need to step one down (arr[r] > pivot) + if (compare(arr[r], pivot) > 0) { + --r; + } + + // the r pointer is on the end of the first k elements + if (k <= r) { + to = r; + } else { + from = r + 1; + } + } + + return arr[k]; + } +}; diff --git a/lib/math.js b/lib/math.js index c937a9848..dad929236 100644 --- a/lib/math.js +++ b/lib/math.js @@ -370,6 +370,7 @@ function create (config) { require('./function/utils/format')(math, _config); require('./function/utils/import')(math, _config); require('./function/utils/map')(math, _config); + require('./function/utils/partitionSelect')(math, _config); require('./function/utils/print')(math, _config); require('./function/utils/sort')(math, _config); require('./function/utils/typeof')(math, _config); diff --git a/test/function/utils/partitionSelect.test.js b/test/function/utils/partitionSelect.test.js new file mode 100644 index 000000000..3322e9acb --- /dev/null +++ b/test/function/utils/partitionSelect.test.js @@ -0,0 +1,109 @@ +var assert = require('assert'); +var error = require('../../../lib/error/index'); +var math = require('../../../index'); +var matrix = math.matrix; +var partitionSelect = math.partitionSelect; + +describe('partitionSelect', function() { + + it('should sort an array with numbers', function() { + assert.equal(partitionSelect([5,10,1], 0), 1); + assert.equal(partitionSelect([5,10,1], 1), 5); + assert.equal(partitionSelect([5,10,1], 2), 10); + }); + + it('should sort an array with strings', function() { + assert.equal(partitionSelect(['C', 'B', 'A', 'D'], 0), 'A'); + assert.equal(partitionSelect(['C', 'B', 'A', 'D'], 1), 'B'); + assert.equal(partitionSelect(['C', 'B', 'A', 'D'], 2), 'C'); + assert.equal(partitionSelect(['C', 'B', 'A', 'D'], 3), 'D'); + }); + + it('should sort a Matrix', function() { + assert.equal(partitionSelect(matrix([5,10,1]), 0), 1); + assert.equal(partitionSelect(matrix([5,10,1]), 1), 5); + assert.equal(partitionSelect(matrix([5,10,1]), 2), 10); + }); + + it('should sort an array in ascending order', function() { + assert.equal(partitionSelect([5,10,1], 0, 'asc'), 1); + assert.equal(partitionSelect([5,10,1], 1, 'asc'), 5); + assert.equal(partitionSelect([5,10,1], 2, 'asc'), 10); + }); + + it('should sort an array in descending order', function() { + assert.equal(partitionSelect([5,10,1], 0, 'desc'), 10); + assert.equal(partitionSelect([5,10,1], 1, 'desc'), 5); + assert.equal(partitionSelect([5,10,1], 2, 'desc'), 1); + }); + + it('should sort an array with a custom compare function', function() { + function sortByLength (a, b) { + return a.length - b.length; + } + assert.equal(partitionSelect(['Langdon', 'Tom', 'Sara'], 0, sortByLength), 'Tom'); + assert.equal(partitionSelect(['Langdon', 'Tom', 'Sara'], 1, sortByLength), 'Sara'); + assert.equal(partitionSelect(['Langdon', 'Tom', 'Sara'], 2, sortByLength), 'Langdon'); + }); + + it('should mutate the input array, leaving it partitioned at k', function() { + var arr = [3, 2, 4, 6, -2, 5]; + partitionSelect(arr, 3); + + for (var i = 0; i < 3; ++i) { + assert.ok(arr[i] <= arr[3]); + } + assert.ok(arr[3] === 4); + for (var i = 4; i < arr.length; ++i) { + assert.ok(arr[3] <= arr[i]); + } + }); + + it('should mutate the input matrix, leaving it partitioned at k', function() { + var m = matrix([3, 2, 4, 6, -2, 5]); + partitionSelect(m, 3); + + m.forEach(function (value, index, matrix) { + if (index[0] < 3) { + assert.ok(value <= 4); + } else if (index[0] === 3) { + assert.ok(value === 4); + } else { + assert.ok(value >= 4); + } + }); + }); + + it('should throw an error if called with a multi dimensional matrix', function() { + assert.throws(function() { partitionSelect(matrix([[1,2],[3,4]]), 2) }, /Only one dimensional matrices supported/); + }); + + it('should throw an error if called with a non-negative k, within the bounds of the array', function() { + assert.throws(function() { partitionSelect([1], -2) }, /k must be a non-negative integer/); + assert.throws(function() { partitionSelect([3, 2, 1], 1.2) }, /k must be a non-negative integer/); + assert.throws(function() { partitionSelect([3, 2, 1], 3) }, /k out of bounds/); + assert.throws(function() { partitionSelect([], 0) }, /k out of bounds/); + }); + + it('should throw an error if called with unsupported type', function() { + assert.throws(function() { partitionSelect(2, 2) }, math.error.UnsupportedTypeError); + assert.throws(function() { partitionSelect('string', 2) }, math.error.UnsupportedTypeError); + assert.throws(function() { partitionSelect([1], 'string') }, math.error.UnsupportedTypeError); + assert.throws(function() { partitionSelect([1], 1, 'function') }, math.error.UnsupportedTypeError); + assert.throws(function() { partitionSelect([1], 1, {}) }, math.error.UnsupportedTypeError); + }); + + it('should throw an error if called with invalid number of arguments', function() { + assert.throws(function() { partitionSelect() }, math.error.ArgumentsError); + assert.throws(function() { partitionSelect([]) }, math.error.ArgumentsError); + assert.throws(function() { partitionSelect([], 2, 'foo', 3) }, math.error.ArgumentsError); + }); + + /* + it('should LaTeX sort', function () { + var expression = math.parse('sort([3,2,1])'); + assert.equal(expression.toTex(), '\\mathrm{sort}\\left(\\begin{bmatrix}3\\\\2\\\\1\\\\\\end{bmatrix}\\right)'); + }); + */ + +});