查找数组中的第 K 个最大元素

写在前面

本文会从一个小问题出发,寻求一题多解,不会对涉及到的算法有详尽的解释。所以在阅读本文前,期望读者已有一定的算法基础,如排序算法、二分查找等。

问题

描述

设数组 A[0..N-1] 存在 N 个无序整数,找到数组 A 中的第 K(1≤K≤N) 大数。
注意:结果是顺序排序后的第 K 个最大的元素,而不是第 K 个不同的最大元素。

示例

输入:A=[3, 2, 3, 1, 2, 4, 5, 5, 6], K=3
输出:5

解题思路

Array.prototype.sort()

利用 JavaScript 标准库中提供的 sort() 方法,对数组元素进行从大到小排序。
由于不同的 JavaScript 引擎sort() 方法的具体实现不同,因此无法保证排序的时间和空间复杂度。

function findKthLargest(nums, k) {
  nums.sort((a, b) => b - a)
  return nums[k - 1]
}

选择排序

利用我们生活中最直观的查找方式:每次从数组剩余元素中找到最大的元素,并将其从数组中剔除,直至进行到第 K 次操作。时间复杂度为 O(n^2)

let i = 0
while (i <= k) {
  if (!nums.length) break
  let j = 0
  for (let index = 0; index < nums.length; index++) {
    if (nums[index] > nums[j]) {
      j = index
    }
  }
  if (i === k - 1) return nums[j]
  nums.splice(j, 1)
  i += 1
}

这与选择排序原理非常相似,不同之处在于,选择排序过程中没有将每次查找到的最大元素从数组中剔除,而是把它添加到已排序序列 A[0..i-1] 的末尾。

for (let i = 0; i < nums.length - 1; i++) {
  let max = i
  for (let j = i + 1; j < nums.length; j++) {
    if (nums[j] > nums[max]) {
      max = j
    }
  }
  [nums[i], nums[max]] = [nums[max], nums[i]]
}

快速排序


以下这种快速排序的写法,是从 Free Pascal 的 Demo 目录文件中学到的,也一直沿用至今。
平均时间复杂度为 O(nlogn),为了避免快速排序退化至 O(n^2),采用了随机化划分基准。

function qsort(l, r) {
  let i = l
  let j = r
  // const x = Math.floor((i + j) / 2)
  const x = Math.floor(Math.random() * (j - i + 1)) + i
  const pivot = nums[x]
  while (i <= j) {
    while (nums[i] > pivot) i += 1
    while (pivot > nums[j]) j -= 1
    if (i <= j) {
      [nums[i], nums[j]] = [nums[j], nums[i]]
      i += 1
      j -= 1
    }
  }
  if (l < j) qsort(l, j)
  if (i < r) qsort(i, r)
}

优化快速排序

根据题目要求,仅需取第 K 大数即可,但在使用快速排序的过程中,把 N 个元素都进行了排序。回顾一下快速排序的思想,在分割操作时,选定划分的基准,比基准大的元素置于基准的左侧;比基准小的元素置于基准的右侧。
为了减少不必要的递归,可以将基准所在的下标与 K 进行比较,从而确定需要递归的区间,缩小递归范围。时间复杂度为 O(n)
设当前的基准 p,将 A 数组经过调整后划分为 S(Si≥p)和 T(Tj<p),此时将有如下两种情况:

  • 如果 S 中的元素个数小于 K,那么第 K 大数存在于 T 中。 T 中的 K - length(S) 大的数即为所求;
  • 如果 S 中的元素个数大于等于 K,那么第 K 大数存在于 S 中。
function search(l, r, kth) {
  if (l === r) return nums[l]
  let i = l
  let j = r
  const x = Math.floor(Math.random() * (j - i + 1)) + i
  const pivot = nums[x]
  while (i <= j) {
    while (nums[i] > pivot) i += 1
    while (pivot > nums[j]) j -= 1
    if (i <= j) {
      [nums[i], nums[j]] = [nums[j], nums[i]]
      i += 1
      j -= 1
    }
  }
  const n = j - l + 1
  if (kth <= n) return search(l, j, kth)
  return search(j + 1, r, kth - n)
}

堆排序

将 N 个元素插入大根中,完成建堆过程。每次从堆顶取出最大元素后,将堆进行向下调整,保持大根堆性质,直至取到第 K 大元素。时间复杂度为 O(nlogn)

const heap = []
let size = -1

function siftup(pos) {
  let child = pos
  while (child) {
    const father = Math.floor((child - 1) / 2)
    if (heap[father] < heap[child]) {
      [heap[father], heap[child]] = [heap[child], heap[father]]
    } else {
      break
    }
    child = father
  }
}

function siftdown(pos) {
  let father = pos
  let child = father * 2 + 1
  while (child <= size) {
    if (child + 1 <= size && heap[child] < heap[child + 1]) {
      child += 1
    }
    if (heap[father] < heap[child]) {
      [heap[father], heap[child]] = [heap[child], heap[father]]
    } else {
      break
    }
    father = child
    child = father * 2 + 1
  }
}

function insert(x) {
  size += 1
  heap[size] = x
  siftup(size)
}

function getMax() {
  return heap[0]
}

function removeRoot() {
  if (size < 0) return;
  [heap[0], heap[size]] = [heap[size], heap[0]]
  size -= 1
  siftdown(0)
}

function build() {
  for (let i = 0; i < nums.length; i++) {
    insert(nums[i])
  }
}

次数统计

统计 A 数组中元素 Ai 出现的次数,并记录出现的最小值和最大值。线性的枚举区间 [max, min],累加当前枚举到的数字 x 在 A 数组中出现的次数,直至累加结果大于或等于 K 时,当前的 x 即为第 K 大数。时间复杂度为 O(n)

const count = {}
nums.forEach((x) => {
  count[x] = count[x] ? count[x] + 1 : 1
})
let sum = 0
for (let i = max; i >= min; i -= 1) {
  sum += count[i] || 0
  if (sum >= k) return i
}

二分答案

利用逆向思维,将原来的从问题寻找答案,转变成假设答案是什么,进而验证答案的正确性。
由于答案一定位于由 A 数组中的最大值和最小值组成的 [min, max] 区间内,且该区间满足区间单调性,所以可以通过对该答案区间进行二分查找,验证答案是否符合条件,不断地缩短区间,最终得出答案。时间复杂度为 O(nlogn)
这里的验证答案,即是验证数组 A 中有多少个元素大于或等于当前枚举的答案。

function count(x) {
  return nums.reduce((counter, i) => (i >= x ? counter + 1 : counter), 0)
}

function search() {
  let left = Math.min(...nums)
  let right = Math.max(...nums)
  while (left <= right) {
    const mid = Math.floor((left + right) / 2)
    const sum = count(mid)
    if (sum >= k) {
      left = mid + 1
    } else {
      right = mid - 1
    }
  }
  return right
}

参考

Kth Largest Element in an Array - LeetCode

Show Comments