Merge pull request #371 from BigFav/quickselect

Implemented a partition-based selection function

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];
+  }
+};

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);

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)');
+  });
+  */
+
+});