Leetcode 刷題 pattern - Two Heaps


Posted by Po-Jen on 2020-03-15

前言

身在大 CS 時代,有越來越多人投入刷題的行列,在眼花撩亂的題海中,要想有效率地刷題,就需要掌握題目解法中,可以在許多地方應用的觀念,才能以簡禦繁。 Huli 寫的 程式解題新手入門注意事項當我們在學程式時,要學的到底是什麼? 講得很好,寫題目是為了學會解題的思考法,確保自己掌握重要的資料結構跟演算法。這也是我寫這系列文章的原因,我想把多個散落在各處的題目銜接起來,以後看到相似的問題就可以舉一反三,而不是去背各題目的解法。

另外,這系列文章提供的做法雖然可以讓面試準備事半功倍,但不一定適合你。如果你比較喜歡透過刷很多題來鍛鍊演算法跟資料結構的直覺,那你可以參考 這篇短文,踏上刷遍一千題、把 computational thinking 刻進腦袋中的道路。

接下來,就讓我們進入今天的正題 - Two Heaps。

Two Heaps 簡介

Two Heaps 顧名思義就是要用到兩個 heap,如果你不知道什麼是 heap,那請你先去學習一下,這篇文章 還不錯可以看看。

了解什麼是 heap 之後,再來猜猜看什麼是 Two Heaps,如果只是用兩個 max heap (或兩個 min heap) 好像沒什麼意思,但如果用一個 min heap,跟一個 max heap,好像就有點搞頭喔!例如說我有 10 個 integer - 0, 1, 2, ..., 9,我可以把 0~4 都放進 max heap,這時 max heap 的頂端會是 4;然後把 5~9 放進 min heap,這時 min heap 的頂端會是 5,這時候如果我要算 median,就只要把 max heap 的頂端跟 min heap 的頂端相加除以 2,就能得到答案。這就是 Two Heaps 的一個妙用!

讓我們直接來看個範例,練習一下。

Two Heaps 的第一個範例 - Leetcode #295 - Find Median from Data Stream

題目

img

暴力解

這題最直觀的暴力解,就是把所有看過的資料都存到一個 array 裡,每次只要 addNum() 被呼叫加入新 data,就把 array sort 一遍;而當 findMedian() 被呼叫,就取 array 中間的值回傳(如果是偶數個就取兩個平均)。

這種做法很簡單,但可想而知效率不是很好,因為每次加入新 data,都要把已有的 data 再重新 sort 一遍,所以 addNum() 的時間複雜度是 O(nlogn),findMedian() 的時間複雜度是 O(1)。

套一句我演算法老師常問的問題 - "Can we do better?",答案必然是,"Yes! Professor"。

Two Heaps 解法

用 Two Heaps 的話,我們可以先創造一個 min heap 跟一個 max heap,在 addNum() 被呼叫的時候,盡量保持 min heap 跟 max heap 的大小一致(大小最多只能差 1),這樣在 findMedian() 的時候,只要從 heap 的頂端取值就好。以時間複雜度來說,addNum() 降為 O(logn)(如果你不知道為什麼,去複習一下 heapify!),findMedian() 還是 O(1)。

假設我們讓 max heap 可以比 min heap 多一個 element(你也可以讓 min heap 多一個 element,只有實作上的小差異),那可以過一個簡單的範例:

img

以下是我的 C++ 實作(Leetcode discussion 中有更簡潔的實作,但我放上來的都是我覺得比可讀性較佳,也比較容易在面試中想出的做法):

class MedianFinder {
public:
    /** initialize your data structure here. */
    MedianFinder() {
    }

    void addNum(int num) {
        if(maxHeap.empty() or num <= maxHeap.top()) {
            maxHeap.push(num);
        }
        else {
            minHeap.push(num);
        }

        if( maxHeap.size() > minHeap.size() + 1) {
            minHeap.push(maxHeap.top());
            maxHeap.pop();
        }
        else if( minHeap.size() > maxHeap.size() ) {
            maxHeap.push(minHeap.top());
            minHeap.pop();
        }
    }

    double findMedian() {
        if(maxHeap.size() == minHeap.size()) {
            return maxHeap.top() / 2.0 + minHeap.top() / 2.0;
        }

        return maxHeap.top();
    }

private:
    priority_queue<int, vector<int>, greater<int>> minHeap;
    priority_queue<int> maxHeap; // keep one more element in maxHeap if there are odd number of elements
};

Two Heaps 的第二個範例 - Leetcode #480 - Sliding Window Median

題目

img

又是一題要求 Median 的,形跡可疑,看來又是要用 Two Heaps 了。

Two Heaps 解法

暴力姐大家大概都想得到,我就不贅述了,直切重點。一開始,可能會想說,哼,還不就是要用 Two Heaps,讓我先來宣告兩個 heap,不過,當你仔細看看題目,會發現當 sliding window 移動,你得從 heap 中刪除元素,這 ...,如果元素不是在 heap 的頂端,好像有點麻煩啊!

在實作上,會遇到好幾個問題(都是我在 Leetcode submit 時的血淚教訓 ^ ^):

  1. 可能有重複元素,所以也不能只用 set,要用 multiset,另外刪除時要刪除 iterator,不要刪除 value,不然可能會誤刪好幾個 element
  2. 縮減完 window 後才平衡兩個 heap 的 element 數,不然會不平衡
  3. 算 median 時注意 int overflow

這題還結合了 sliding window,所以實作出來的 code 比較長!如果你已經開始瑟瑟發抖,建議去快速看一下延伸閱讀裡面的 sliding window 篇,熟悉一下 sliding window 的實作再回來,code 就會變得親切許多。

概念上沒什麼好多說的了,直接把手弄髒,你就會學習了:

class Solution {
public:
    vector<double> medianSlidingWindow(vector<int>& nums, int k) {
        vector<double> medians;
        int windowStart = 0, n = nums.size();

        for(int windowEnd = 0; windowEnd < n; windowEnd++) {
            // Expand windowEnd and add to heaps
            if(maxHeap.empty() or nums[windowEnd] <= *maxHeap.rbegin()) {
                maxHeap.insert(nums[windowEnd]);
            }
            else {
                minHeap.insert(nums[windowEnd]);
            }

            // If window size too big, move windowStart to reduce window
            int windowSize = maxHeap.size() + minHeap.size();
            if(windowSize > k) {
                int toDelete = nums[windowStart++];
                if(maxHeap.count(toDelete)) {
                    maxHeap.erase(maxHeap.find(toDelete));
                }
                else {
                    minHeap.erase(minHeap.find(toDelete));
                }
            }

            // Balance maxHeap and minHeap
            if(maxHeap.size() > minHeap.size() + 1) {
                minHeap.insert(*maxHeap.rbegin());

                auto it = maxHeap.end(); 
                --it;
                maxHeap.erase(it);
            }
            else if(minHeap.size() > maxHeap.size()) {
                maxHeap.insert(*minHeap.begin());
                minHeap.erase(minHeap.begin());
            }

            // Get median and put into medians
            windowSize = maxHeap.size() + minHeap.size();
            if(windowSize == k) {
                if(k % 2 == 0) {
                    long long sum = *minHeap.begin() + *maxHeap.rbegin();
                    medians.push_back( sum / 2.0 );
                }
                else {
                    medians.push_back( *maxHeap.rbegin() );
                }
            }    
        }

        return medians;
    }

private:
    multiset<int> maxHeap; // use rbegin() to get max value
    multiset<int> minHeap; // use begin() to get min value
};

Two Heaps 的第三個範例 - Leetcode #436 - Find Right Interval

題目

img
img

Two Heaps 解法

事先聲明一下,這題的 Two Heaps 解法寫起來又臭又長,leetcode discussion 上有其他簡潔有效率的解法,所以如果你真的被考到這題,不一定要用這個解法。但因為這題可以用 Two Heaps 解,而且概念上也還算巧妙,斟酌之後還是決定放上來給大家參考,擴展一下思路。

先講一下概念:

我們可以把所有的 interval 放進兩個 heap,maxStartHeap 的頂端會是開始時間最大的 interval;maxEndHeap 的頂端則是結束時間最大的 interval。接下來,我們就坐這幾件事:

  1. 從 maxEndHeap 的頂端取出 interval,因為這個 interval 的結束時間最晚,我們把它稱為 topEnd。
  2. 從 maxStartHeap 的頂端開始取出 interval,並跟 topEnd 比較,直到這個 interval 的開始時間最接近 topEnd 的結束時間(但仍在 topEnd 的右邊),我們稱這個 interval 為topStart。(因為 topEnd 已經是結束時間最晚的 interval,所以比 topStart 還晚開始的 interval,已經不可能在答案中出現,就不需要繼續存放在 maxStartHeap 中)
  3. 把 topStart 的 index 放進答案中,如果找不到就是 -1。
  4. 把 topStart 放回 maxStartHeap,因為他仍然可能是下一個 topEnd 的 topStart。
  5. 重複步驟 1-4 直到 maxEndHeap 是空的。

實作比較複雜:

class Solution {
public:
    vector<int> findRightInterval(vector<vector<int>>& intervals) {
        int n = intervals.size();

        // heap for finding the maximum start, we also need to store the index for putting the ans correctly
        priority_queue<pair<pair<int, int>, int>, vector<pair<pair<int, int>, int>>, startCompare> maxStartHeap;
        // heap for finding the maximum end
        priority_queue<pair<pair<int, int>, int>, vector<pair<pair<int, int>, int>>, endCompare> maxEndHeap;

        vector<int> result(n, -1);
        for (int i = 0; i < n; ++i) {
            maxStartHeap.push(make_pair(make_pair(intervals[i][0], intervals[i][1]), i));
            maxEndHeap.push(make_pair(make_pair(intervals[i][0], intervals[i][1]), i));
        }

        // go through all the intervals to find each interval's next interval
        for (int i = 0; i < n; i++) {
            // let's find the next interval of the interval which has the highest 'end'
            auto topEnd = maxEndHeap.top();
            maxEndHeap.pop();

            if (maxStartHeap.top().first.first >= topEnd.first.second) {
                auto topStart = maxStartHeap.top();
                maxStartHeap.pop();

                // find the the interval that has the closest 'start'
                while (!maxStartHeap.empty() && maxStartHeap.top().first.first >= topEnd.first.second) {
                    topStart = maxStartHeap.top();
                    maxStartHeap.pop();
                }

                result[topEnd.second] = topStart.second;

                // put the interval back as it could be the next interval of other intervals
                maxStartHeap.push(topStart);
            }
        }

        return result;
    }

private:
    struct startCompare {
        bool operator() (const pair<pair<int, int>, int> &x, const pair<pair<int, int>, int> &y) {
            return y.first.first > x.first.first;
        }
    };

    struct endCompare {
        bool operator()(const pair<pair<int, int>, int> &x, const pair<pair<int, int>, int> &y) {
            return y.first.second > x.first.second;
        }
    };
};

總結

今天跟大家介紹了一個新的 pattern - Two Heaps,順便提供幾個經典的可應用題,讓大家體驗一下用這個 pattern 的好處。雖然 Two Heaps 的概念很簡單,但如果你對這個 pattern 不熟,實作時要一次寫出 bug-free code 也不是很容易。個人覺得這個 pattern 頗適合用來磨練實作技巧,同時也能展現你靈活運用資料結構 (Heap) 的能力。

延伸閱讀

  1. Leetcode 刷題 pattern - Two Pointer
  2. Leetcode 刷題 pattern - Sliding Window
  3. Leetcode 刷題 pattern - Next Greater Element
  4. Leetcode 刷題 pattern - Fast & Slow Pointer
  5. Leetcode 刷題 pattern - Breadth-First Search
  6. Leetcode 刷題 pattern - Merge Intervals
  7. Leetcode 刷題 pattern - Cyclic Sort

關於作者:
@pojenlai 演算法工程師,對機器人、電腦視覺和人工智慧有少許研究,正致力改善 自己的習慣


#Leetcode #algorithm #Software Engineer









Related Posts

JavaScript 的變數

JavaScript 的變數

[AI人工智能] 一、Python、Anaconda、Pycharm安裝

[AI人工智能] 一、Python、Anaconda、Pycharm安裝

設計模式 - 工廠方法及抽象工廠

設計模式 - 工廠方法及抽象工廠




Newsletter




Comments