【堆】 LeetCode295. 数据流的中位数

题目

解答

#include <iostream> #include <string> #include <list> #include <queue> #include <vector> #include <functional> #include <algorithm> #include <cassert> // 295. Find Median from Data Stream。 // 巧用最大堆和最小堆。 // 将这组数据分为两段，前半段和后半段，后半段的最小值肯定大于前半段的最大值。 // 所以前半段用最大堆，后半段用最小堆。由于堆的性质可以不用每次都将所有数据都排序， // 所以极大地提高了效率。 class MedianFinder { public: /** initialize your data structure here. */ MedianFinder() { } void addNum(int num) { // 前半段和后半段的长度的差值永远小于等于1。 std::size_t max_heap_size = max_heap_.size(); std::size_t min_heap_size = min_heap_.size(); // 如果前半段长度大于后半段，那么就将一个值放到后半段。 if (max_heap_size > min_heap_size) { // 如果前半段的最大值大于要插入的值，那么该插入的值不能放到后半段。 // 因为后半段的值肯定大于前半段的值，所以将其交换再放到后半段。 if (max_heap_.top() > num) { min_heap_.push(max_heap_.top()); max_heap_.pop(); max_heap_.push(num); } else { min_heap_.push(num); } // 如果后半段长度大于前半段，那么就将一个值放到前半段。 } else if (max_heap_size < min_heap_size) { if (min_heap_.top() < num) { max_heap_.push(min_heap_.top()); min_heap_.pop(); min_heap_.push(num); } else { max_heap_.push(num); } } else { if (min_heap_.empty() || min_heap_.top() < num) { min_heap_.push(num); } else { max_heap_.push(num); } } } double findMedian() { std::size_t max_heap_size = max_heap_.size(); std::size_t min_heap_size = min_heap_.size(); if (max_heap_size > min_heap_size) { return max_heap_.top(); } else if (max_heap_size < min_heap_size) { return min_heap_.top(); } else { return (min_heap_.top() + max_heap_.top()) / 2.0; } } private: std::priority_queue<int, std::vector<int>, std::greater<int>> min_heap_; std::priority_queue<int, std::vector<int>, std::less<int>> max_heap_; }; int main() { MedianFinder median_finder; median_finder.addNum(1); std::cout << median_finder.findMedian() << std::endl; median_finder.addNum(2); std::cout << median_finder.findMedian() << std::endl; median_finder.addNum(3); std::cout << median_finder.findMedian() << std::endl; median_finder.addNum(4); std::cout << median_finder.findMedian() << std::endl; median_finder.addNum(5); std::cout << median_finder.findMedian() << std::endl; median_finder.addNum(6); std::cout << median_finder.findMedian() << std::endl; return 0; }