陣列中的第K個最大元素

又是火星人發表於2024-05-28

主管問到這個問題(陣列中的第K個最大元素)。
我首先答了partition演算法時間複雜度O(n)空間複雜度O(1)。主管說不行,要用堆。然後(我回答)用建大根堆,取前k,時間複雜度O(n+klogn)初始化O(n),pop出k個O(klogn),空間複雜度O(k)(使用原陣列建堆,pop出k個)。主管說都不行,然後(主管)給了個小根堆的演算法,逐個把數插入堆中。維持小根堆大小為k,如果超出就出堆。這樣的話,相當於把n-k個小的pop出去了剩下的小根堆就存了k個最大的。時間複雜度O(nlogk)空間複雜度O(k)
主管預設了場景了,說是資料規模很大,但是隻要前幾個,也就是說n很大k很小。所以說nlogk是不可能比得過n+klogn的。哪一個都沒我的演算法好啊。無語。

主管給出的要求通常是希望看到你對演算法和資料結構的全面理解以及實際場景中的應用。針對你提到的查詢前 k 大元素的問題,主管希望你使用小根堆,這是因為它在處理某些特定場景(如資料規模非常大但只需前幾大元素)時具有顯著優勢。以下是對這幾種方法的詳細分析,以及為什麼在這種場景下使用小根堆是一個更優的選擇。

常見解法及其時間和空間複雜度

1. Partition Algorithm(基於快速選擇的演算法)

  • 時間複雜度: 平均 O(n),最壞 O(n^2)
  • 空間複雜度: O(1)

這種方法在實際應用中可能會遇到最壞情況(例如,選擇的基準總是最大或最小元素),導致效能不穩定。

2. 大根堆(Max-Heap)

  • 構建堆的時間複雜度: O(n)
  • 取前 k 大元素的時間複雜度: O(k log n)
  • 總體時間複雜度: O(n + k log n)
  • 空間複雜度: O(1)(使用原陣列)

這種方法在資料量很大時,構建堆的時間複雜度是 O(n),而取前 k 大元素的時間複雜度是 O(k log n)。雖然總體時間複雜度比較好,但對資料量極大而 k 很小時,仍然可能不夠高效。

3. 小根堆(Min-Heap)

  • 時間複雜度: O(n log k)
  • 空間複雜度: O(k)

這種方法的思路是維持一個大小為 k 的小根堆,當堆的大小超過 k 時,將堆頂元素(最小的元素)移除。這樣,最終堆中儲存的就是前 k 大的元素。

為什麼小根堆在這種場景下更優?

在這種場景下,n 很大而 k 很小,使用小根堆的時間複雜度是 O(n log k),這比大根堆的 O(n + k log n) 更有優勢。具體來說:

  • 插入操作: 每次插入一個元素到堆中,操作時間是 O(log k),而總共需要插入 n 次,總時間是 O(n log k)。
  • 堆的大小: 由於堆的大小始終維持在 k,因此空間複雜度是 O(k)。

小根堆實現程式碼示例

import java.util.PriorityQueue;

class Solution {
    public int findKthLargest(int[] nums, int k) {
        // 建立一個小根堆
        PriorityQueue<Integer> minHeap = new PriorityQueue<>(k);
        
        // 遍歷陣列,將元素插入堆中
        for (int num : nums) {
            if (minHeap.size() < k) {
                // 如果堆的大小小於 k,直接插入
                minHeap.offer(num);
            } else if (num > minHeap.peek()) {
                // 如果堆的大小已經達到 k,並且當前元素大於堆頂元素
                // 先移除堆頂元素,再插入當前元素
                minHeap.poll();
                minHeap.offer(num);
            }
        }
        
        // 堆頂元素就是第 k 大的元素
        return minHeap.peek();
    }
}

總結

  • Partition Algorithm 適用於一般場景,但在極端情況下效能不穩定。
  • 大根堆 在需要多次取最大值時非常有效,但在處理極大資料集而只需少量最大值時,時間複雜度不如小根堆。
  • 小根堆n 很大而 k 很小的情況下,具有更優的時間複雜度 O(n log k) 和適度的空間複雜度 O(k),因此在這種特定場景下是最優選擇。

主管希望看到你能夠根據具體的場景和需求選擇最合適的演算法。理解並靈活運用不同的資料結構和演算法,是展示你能力的關鍵。

相關文章