前言
身在大 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
題目
暴力解
這題最直觀的暴力解,就是把所有看過的資料都存到一個 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,只有實作上的小差異),那可以過一個簡單的範例:
以下是我的 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
題目
又是一題要求 Median 的,形跡可疑,看來又是要用 Two Heaps 了。
Two Heaps 解法
暴力姐大家大概都想得到,我就不贅述了,直切重點。一開始,可能會想說,哼,還不就是要用 Two Heaps,讓我先來宣告兩個 heap,不過,當你仔細看看題目,會發現當 sliding window 移動,你得從 heap 中刪除元素,這 ...,如果元素不是在 heap 的頂端,好像有點麻煩啊!
在實作上,會遇到好幾個問題(都是我在 Leetcode submit 時的血淚教訓 ^ ^):
- 可能有重複元素,所以也不能只用 set,要用 multiset,另外刪除時要刪除 iterator,不要刪除 value,不然可能會誤刪好幾個 element
- 縮減完 window 後才平衡兩個 heap 的 element 數,不然會不平衡
- 算 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
題目
Two Heaps 解法
事先聲明一下,這題的 Two Heaps 解法寫起來又臭又長,leetcode discussion 上有其他簡潔有效率的解法,所以如果你真的被考到這題,不一定要用這個解法。但因為這題可以用 Two Heaps 解,而且概念上也還算巧妙,斟酌之後還是決定放上來給大家參考,擴展一下思路。
先講一下概念:
我們可以把所有的 interval 放進兩個 heap,maxStartHeap 的頂端會是開始時間最大的 interval;maxEndHeap 的頂端則是結束時間最大的 interval。接下來,我們就坐這幾件事:
- 從 maxEndHeap 的頂端取出 interval,因為這個 interval 的結束時間最晚,我們把它稱為 topEnd。
- 從 maxStartHeap 的頂端開始取出 interval,並跟 topEnd 比較,直到這個 interval 的開始時間最接近 topEnd 的結束時間(但仍在 topEnd 的右邊),我們稱這個 interval 為topStart。(因為 topEnd 已經是結束時間最晚的 interval,所以比 topStart 還晚開始的 interval,已經不可能在答案中出現,就不需要繼續存放在 maxStartHeap 中)
- 把 topStart 的 index 放進答案中,如果找不到就是 -1。
- 把 topStart 放回 maxStartHeap,因為他仍然可能是下一個 topEnd 的 topStart。
- 重複步驟 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) 的能力。