Function select

tfjs-backend-cpu/src/kernels/TopK_impl.ts:43–94

* Partitions array where all elements smaller than the (k+1) smallest element * are found to the left of it, and all larger to the right of it. * Based on the Floyd-Rivest Algorithm, ref: * https://en.wikipedia.org/wiki/Floyd%E2%80%93Rivest_algorithm * @param array: Array to partition * @param

(array: Pair[], k: number, left = 0, right = array.length - 1)

 *           when left = 0
 */
function select(array: Pair[], k: number, left = 0, right = array.length - 1) {
  while (right > left) {
    // Use select recursively to sample a smaller set of size s
    // the arbitrary constants 600 and 0.5 are used in the original
    // version to minimize execution time.
    if (right - left > 600) {
      const n = right - left + 1;
      const i = k - left + 1;
      const z = Math.log(n);
      const s = 0.5 * Math.exp(2 * z / 3);
      const sd = 0.5 * Math.sqrt(z * s * (n - s) / n) * Math.sign(i - n / 2);
      const newLeft = Math.max(left, Math.floor(k - i * s / n + sd));
      const newRight = Math.min(right, Math.floor(k + (n - i) * s / n + sd));
      select(array, k, newLeft, newRight);
    }
    // partition the elements between left and right around t
    const t = array[k];
    let i = left;
    let j = right;

    util.swap(array, left, k);

    if (comparePair(array[right], t) > 0) {
      util.swap(array, left, right);
    }
    while (i < j) {
      util.swap(array, i, j);
      i++;
      j--;
      while (comparePair(array[i], t) < 0) {
        i = i + 1;
      }
      while (comparePair(array[j], t) > 0) {
        j = j - 1;
      }
    }
    if (comparePair(array[left], t) === 0) {
      util.swap(array, left, j);
    } else {
      j = j + 1;
      util.swap(array, j, right);
    }
    // Adjust left and right towards the boundaries of the subset
    // containing the (k - left + 1)th smallest element.
    if (j <= k) {
      left = j + 1;
    }
    if (k <= j) {
      right = j - 1;
    }
  }
}

export function topKImpl<T extends Tensor, R extends Rank>(
    x: TypedArray, xShape: number[], xDtype: NumericDataType, k: number,