[譯] Swift 算法學院 - 查找數組中第 K 大值

本篇是來自 Swift 算法學院的翻譯的一篇文章，Swift 算法學院 致力於使用 Swift 實現各類算法，對想學習算法或者複習算法的同窗很是有幫助，講解思路很是清楚，每一篇都有詳細的例子解釋。 更多翻譯的文章還能夠查看這裏。

第K大元素

給定一個數組 a ，寫一個算法找出第K大的元素。

好比在 第一大 的元素是最大元素。若是數組有 n 個元素，第 n 大 元素爲最小值，中間最大爲 第n/2 大值。

原始方案

下面的算法是半原生的。它的時間複雜度是 O(n log n)，由於它須要先排序，所以須要額外的 O(n) 的空間。

func kthLargest(a: [Int], k: Int) -> Int? {
  let len = a.count
  if k > 0 && k <= len {
    let sorted = a.sorted()
    return sorted[len - k]
  } else {
    return nil
  }
}
複製代碼

kthLargest() 函數有兩個參數，整數型數組 a 和 k 用來表示第 k 大的元素。

舉例說明一下這個算法的原理，假定 k = 4 ， 數組以下：

[ 7, 92, 23, 9, -1, 0, 11, 6 ]
複製代碼

最開始沒法直接找到第 k 大的元素，可是排序後就很是簡單了，排序後以下：

[ -1, 0, 6, 7, 9, 11, 23, 92 ]
複製代碼

如今只須要取 a.count - k 對應的值：

a[a.count - k] = a[8 - 4] = a[4] = 9
複製代碼

固然若是須要找 第 k 小 的值時用 a[k-1] 便可

更快的算法

這個算法借鑑了 二分查找 和 快排 的思想，時間複雜度爲 O(n)

不斷調用二分查找將數組分割成一半又一半，快速的縮小查詢的值的範圍。

快速排序也分割數組，把小於軸值的移至左邊，全部大於軸值的移至右邊。通過某個軸值分區後，軸值所在的位置就是排序後最終位置。能夠利用這一點來提升算法。

下面介紹如何工做：隨機選一個值做爲軸值進行分區，像二分查找同樣繼續對左右分區進行處理，直到剛好一個軸值是在 k-th 位置。

舉個例子說明一下，在下面的數組中找 第 4 大的元素：

[ 7, 92, 23, 9, -1, 0, 11, 6 ]
複製代碼

該算法對查找第k小值也是很簡單的，來讓咱們試試查找k = 4 的最小值。

咱們不用先對數組排序，隨機選一個值好比 11 做爲軸值進行分區，結果以下：

[ 7, 9, -1, 0, 6, 11, 92, 23 ]
 <------ smaller    larger -->
複製代碼

根據結果，比 11 小的值在左邊，大的值在右邊。11 在它的最終位置上，索引值爲 5 ， 所以第 4 小的值確定是在左邊的位置能夠忽略其餘的部分：

[ 7, 9, -1, 0, 6, x, x, x ]
複製代碼

再隨機選一個軸值好比 6 將數組分區，結果以下：

[ -1, 0, 6, 9, 7, x, x, x ]
複製代碼

軸值 6 的索引值爲 2，顯然第 4 大的值在右邊分區，能夠忽略左邊的分區了：

[ x, x, x, 9, 7, x, x, x ]
複製代碼

重複以上操做後以下：

[ x, x, x, 7, 9, x, x, x ]
複製代碼

軸值 9 的索引值爲 4，並且這正是要查找的！能夠看到咱們不須要對數組排序，用不多的步數就能實現。

實現方法以下：

public func randomizedSelect<T: Comparable>(_ array: [T], order k: Int) -> T {
  var a = array

  func randomPivot<T: Comparable>(_ a: inout [T], _ low: Int, _ high: Int) -> T {
    let pivotIndex = random(min: low, max: high)
    a.swapAt(pivotIndex, high)
    return a[high]
  }

  func randomizedPartition<T: Comparable>(_ a: inout [T], _ low: Int, _ high: Int) -> Int {
    let pivot = randomPivot(&a, low, high)
    var i = low
    for j in low..<high {
      if a[j] <= pivot {
        a.swapAt(i, j)
        i += 1
      }
    }
    a.swapAt(i, high)
    return i
  }

  func randomizedSelect<T: Comparable>(_ a: inout [T], _ low: Int, _ high: Int, _ k: Int) -> T {
    if low < high {
      let p = randomizedPartition(&a, low, high)
      if k == p {
        return a[p]
      } else if k < p {
        return randomizedSelect(&a, low, p - 1, k)
      } else {
        return randomizedSelect(&a, p + 1, high, k)
      }
    } else {
      return a[low]
    }
  }

  precondition(a.count > 0)
  return randomizedSelect(&a, 0, a.count - 1, k)
}
複製代碼

爲了提升可讀性，這個函數分紅三個內部函數：

• randomPivot() 隨機選取一個數字，而後放在當前分區的最後一個位置（這是Lomuto 分區方式所規定的，更多介紹請看快排）
• randomizedPartition() 是快排中 Lomuto 分區方法。當完成後，隨機軸值在的位置就是排序後的最終位置。返回軸值所在的位置。
• randomizedSelect() 作全部的髒活累活。先調用分區函數，後決定再作什麼。若是軸值索引值等於 k ,那麼該值正是查找值，完成查找。若是 k 比該索引值小，那麼查找值必定在左邊分區，遞歸調用就能夠了，不然就確定是在右邊分區中。

很是😎，是否是？ 快排的指望複雜度爲 o(n log n)， 可是由於只把數組分紅愈來愈小的分區，randomizedSelect() 的時間複雜度爲 O(n)。

注意：該函數式計算數組中 第k 小元素，k 是從 0 開始的。若是須要 第k 大元素，應調用 a.count - k。

做者 Daniel Speiser 修改 Matthijs Hollemans 譯者KeithMorning