JavaScript BFPRT 算法

BFPRT 算法

背景

在一组数中求其前 k 小的数，简称TOP-K问题。而目前解决TOP-K问题最有效的算法即是BFPRT算法，其又称为中位数的中位数算法，该算法由Blum、Floyd、Pratt、Rivest、Tarjan提出，最坏时间复杂度为 O(n)

目的

找到数组中第 k 小的数

步骤

跟快排思路相似，划分为小于区、等于区、大于区三段区域，检查等于区边界 L - R 是否命中 k，如果没命中再去大于区或小于区重复查找。
BFPRT 算法解决了如何选择划分值 P

1 找划分值
2 根据划分值把数组划分为小于区、等于区、大于区三段区域
3 当传入位置 i 命中等于区, 那么arr[i] 就是我们找到的数

划分值思路

1 将 N 个元素已 5 个一组进行划分，多余的再分一组。 复杂度O(1)，N / 5 组

2 对每组 5 个数进行排序查找中位数�O(n)
3 从每组的中位数中找新的中位数，也就是整个数组的中位数 N / 10 确定了划分值 pivot

4 根据划分值 pivot 对数组进行大于区、小于区、等于区划分，确定 等于区 L - R 的位置

5 当传入的位置 i (首次为k-1)在 L - R 之中，那么此位置 arr[i] 就是要找到数
6 递归的确定数组的划分值直到确定初次传入的位置的等于区, 找到此数

核心代码

function select(arr, begin, end, i) {
  if (begin === end)
    return arr[begin]
  const pivot = medianOfMedians(arr, begin, end) // 中位数的中位数，确定划分值
  const pivotRange = partition(arr, begin, end, pivot) // 等于区划分位置
  
  if (i >= pivotRange[0] && i <= pivotRange[1])
    return arr[i] // 在等于区内 找到此位置的数
  else if (i < pivotRange[0])
    return select(arr, begin, pivotRange[0] - 1, i)
  else 
    return select(arr, pivotRange[1] + 1, end, i)
}

function medianOfMedians(arr, begin, end) {
  const num = end - begin + 1
  const offset = num % 5 === 0 ? 0 : 1
  // 有多少组，5 个数一组, 多的 + 1组
  let mArr = new Array((num / 5 + offset) | 0)
  
  for (let i = 0; i < mArr.length; i++) {
    // 从每组的位置开始
    const beginI = begin + i * 5
    const endI = beginI + 4
    
    // 返回每组的中位数
    mArr[i] = getMedian(arr, beginI, Math.min(end, endI))
  }
  
  // 返回中位数的中位数, 也就是交集的中位数。
  return select(mArr, 0, mArr.length - 1, mArr.length / 2)
}

function getMedian(arr, begin, end) {
  insertionSort(arr, begin, end) // 先排序
  const sum = end + begin
  const mid = (sum / 2) + (sum % 2)
  
  return arr[mid] // 返回每小组的中位数
}

BFPRT 全部代码

function getMinKNumsByBFPRT(arr, k) {
  if (k < 1 || k > arr.length) return arr
  const minKth = getMinKthByBFPRT(arr, k) // 找到数组中第 k 小的数
  let res = new Array(k)
  let index = 0
  for (let i = 0; i !== arr.length; i++) {
    if (arr[i] < minKth)
      res[index++] = arr[i]
  }
  for (; index !== res.length; index++) {
    res[index] = minKth
  }
  return res
}

function getMinKthByBFPRT(arr, k) {
  const copyArr = Array.from(arr)
  
  return select(copyArr, 0, copyArr.length - 1, k - 1)
}

function select(arr, begin, end, i) {
  if (begin === end)
    return arr[begin]
  const pivot = medianOfMedians(arr, begin, end) // 中位数的中位数
  const pivotRange = partition(arr, begin, end, pivot) // 等于区划分位置
  
  if (i >= pivotRange[0] && i <= pivotRange[1])
    return arr[i] // 在等于区内 
  else if (i < pivotRange[0])
    return select(arr, begin, pivotRange[0] - 1, i)
  else 
    return select(arr, pivotRange[1] + 1, end, i)
}

function medianOfMedians(arr, begin, end) {
  const num = end - begin + 1
  const offset = num % 5 === 0 ? 0 : 1
  // 有多少组，5 个数一组, 多的 + 1组
  let mArr = new Array((num / 5 + offset) | 0)
  
  for (let i = 0; i < mArr.length; i++) {
    const beginI = begin + i * 5
    const endI = beginI + 4
    
    // 返回每组的中位数
    mArr[i] = getMedian(arr, beginI, Math.min(end, endI))
  }
  
  // 返回中位数的中位数
  return select(mArr, 0, mArr.length - 1, mArr.length / 2)
}

function partition(arr, begin, end, pivotValue) {
  let small = begin - 1
  let cur = begin
  let big = end + 1
  while (cur !== big) {
    if (arr[cur] < pivotValue)
      swap(arr, ++small, cur++)
    else if (arr[cur] > pivotValue)
      swap(arr, cur, --big)
    else
      cur++
  }
  return [ small + 1, big - 1 ]
}

function getMedian(arr, begin, end) {
  insertionSort(arr, begin, end)
  const sum = end + begin
  const mid = (sum / 2) + (sum % 2)
  
  return arr[mid]
}

function insertionSort(arr, begin, end) {
  for (let i = begin + 1; i !== end + 1; i++) {
    for (let j = i; j !== begin; j--) {
      if (arr[j - 1] > arr[j])
        swap(arr, j - 1, j)
      else
        break
    }
  }
}

function swap(arr, n, m) {
  [arr[n], arr[m]] = [arr[m], arr[n]]
}

let arr = [6, 9, 1, 3, 1, 2, 2, 5, 6, 1, 3, 5, 9, 7, 2, 5, 6, 1, 9]

console.log(getMinKNumsByBFPRT(arr, 10))