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

物件屬性描述器 Property Descriptor(一)

物件屬性描述器 Property Descriptor(一)

How to solve the perpetual loading issue in Evernote? Evernote 一直轉圈圈的解決辦法

How to solve the perpetual loading issue in Evernote? Evernote 一直轉圈圈的解決辦法

TypeScript 筆記:推斷、註記與斷言

TypeScript 筆記:推斷、註記與斷言




Newsletter




Comments