[譯] Swift 演算法學院 - 查詢陣列中第 K 大值

KeithSummer發表於2018-02-01
Swift演算法陣列

本篇是來自 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

相關文章