获取数组中最大的m个元素(JavaScript版本)

136 阅读2分钟

一、问题描述

获取数组中最大的m个元素,不需要考虑输出顺序

二、解决方案

1. 利用最小堆

使用数组中的前m个元素构建一个最小堆,此时堆顶元素是这m个元素中最小的元素,然后将后续元素依次跟堆顶元素做比较,如果比堆顶元素小,则说明比该最小堆中的所有元素小,无需入堆,如果比堆顶元素大,则将堆顶元素置换为该元素,并且重新调整堆,则可以获取到最新的最大的m个元素,遍历完成之后最小堆中的元素将为最大的m个元素。

2. 利用快排中的partition函数

调用一次partition函数后,将获取到初次划分的索引,该索引前方的元素都大于该位置的元素,后方的元素都小于该位置的元素,如果此索引整好等于m则前方m个元素就是我们要找的最大的m个元素,否则再两边递归进行查找。

三、源代码

1. 堆排序版本

// 打乱数组顺序
Array.prototype.shuffle = function () {
  let m = this.length, i;
  while (m) {
    i = (Math.random() * m--) >>> 0;
    [this[m], this[i]] = [this[i], this[m]];
  }
  return this;
}

// 调整堆
const adjustHeap = (arr, start, end, cmp = (a, b) => a < b) => {
  let dad = start;
  let son = dad * 2 + 1;
  while (son <= end) {
    if (son + 1 <= end && cmp(arr[son + 1], arr[son])) {
      son++;
    }
    if (cmp(arr[dad], arr[son])) {
      return;
    } else {
      [arr[dad], arr[son]] = [arr[son], arr[dad]];
      dad = son;
      son = dad * 2 + 1;
    }
  }
}

// 建堆
const buildHeap = (arr, len = arr.length) => {
  for (let i = parseInt(len / 2) - 1; i >= 0; i--) {
    adjustHeap(arr, i, len - 1);
  }
}

// 获取数组中最大的m个元素
const getMMax= (arr, m) => {
  buildHeap(arr, m);
  for (let i = m; i < arr.length; i++) {
    if (arr[i] > arr[0]) {
      [arr[0], arr[i]] = [arr[i], arr[0]];
      adjustHeap(arr, 0, m - 1);
    }
  }
  return arr.splice(0, m);
}

// 得到一个0-10的数组并打乱顺序
const arr = Array.from(Array(10), (_, k) => k).shuffle();
console.log(arr);
let res = getNMax(arr, 5);
console.log(res);
[ 8, 2, 5, 7, 0, 9, 6, 4, 1, 3 ]
[ 5, 7, 6, 9, 8 ]

2. 快速排序版本

Array.prototype.shuffle = function () {
  let m = this.length, i;
  while (m) {
    i = (Math.random() * m--) >>> 0;
    [this[m], this[i]] = [this[i], this[m]];
  }
  return this;
}

let partition = (arr, begin, end) => {
  if (begin >= end) {
    return -1;
  }
  let i = begin + 1, j = end;
  while (i < j) {
    if (arr[i] < arr[begin]) {
      [arr[i], arr[j]] = [arr[j], arr[i]];
      j--;
    } else {
      i++;
    }
  }
  (arr[i] <= arr[begin]) && i--;
  [arr[i], arr[begin]] = [arr[begin], arr[i]];
  return i;
}

let QuickSort = (arr, begin, end, m) => {
  let pos = partition(arr, begin, end)
  if (pos < m) {
    return QuickSort(arr, pos + 1, end, m)
  } else if (pos > m) {
    return QuickSort(arr, begin, pos - 1, m)
  } else {
    return arr.slice(0, m)
  }
}
// 获取数组中最大的m个元素
const getMMax = (arr, m) => {
  return QuickSort(arr, 0, arr.length - 1, m)
}

let arr = Array.from(Array(10), (_, k) => k).shuffle()
console.log(arr)
let res = getMMax(arr, 5)
console.log(res)
[ 3, 5, 6, 0, 4, 8, 1, 2, 9, 7 ]
[ 9, 8, 7, 6, 5 ]