Skip to the content.

Ex41 数据流的中位数

题目描述

如何得到一个数据流中的中位数?如果从数据流中读出奇数个数值,那么中位数就是所有数值排序之后位于中间的数值。如果从数据流中读出偶数个数值,那么中位数就是所有数值排序之后中间两个数的平均值。

例如,

[2,3,4] 的中位数是 3

[2,3] 的中位数是 (2 + 3) / 2 = 2.5

解题思路

使用一个最大堆和最小堆,最大堆中存放数据流中小的一部分,最小堆中存放数据流中大的一部分。每一次add新数据的时候,根据两个堆的size,决定插入较小的那一个堆;如果一样的话,就插入小的。在这个基础上进行一次再平衡,即比对最小堆的最小值(堆顶)是否小于最大堆的最大值(堆顶),如果小于的话,交换两者的堆顶元素。

代码

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

  void
  addNum (int num)
  {
    if ((min_heap_.size () + max_heap_.size ()) % 2 == 0)
      {
        min_heap_.push (num);
        if (max_heap_.size () > 0 && min_heap_.top () < max_heap_.top ())
          {
            max_heap_.push (min_heap_.top ());
            min_heap_.pop ();
          }
      }
    else
      {
        if (min_heap_.size () > max_heap_.size ())
          {
            max_heap_.push (num);
            if (min_heap_.size () > 0 && min_heap_.top () < max_heap_.top ())
              {
                int max_top = max_heap_.top (), min_top = min_heap_.top ();
                max_heap_.pop ();
                min_heap_.pop ();
                max_heap_.push (min_top);
                min_heap_.push (max_top);
              }
          }
        else
          {
            min_heap_.push (num);
            if (max_heap_.size () > 0 && min_heap_.top () < max_heap_.top ())
              {
                int max_top = max_heap_.top (), min_top = min_heap_.top ();
                max_heap_.pop ();
                min_heap_.pop ();
                max_heap_.push (min_top);
                min_heap_.push (max_top);
              }
          }
      }
  }

  double
  findMedian ()
  {
    return (max_heap_.size () + min_heap_.size ()) % 2 == 0
               ? (min_heap_.top () + max_heap_.top ()) / 2.0
               : (max_heap_.size () > min_heap_.size () ? max_heap_.top ()
                                                        : min_heap_.top ());
  }

private:
  std::priority_queue<int, std::vector<int>, std::less<int> > max_heap_;
  std::priority_queue<int, std::vector<int>, std::greater<int> > min_heap_;

};

结果

执行结果:通过

执行用时:204 ms, 在所有 C++ 提交中击败了96.70%的用户

内存消耗:41.2 MB, 在所有 C++ 提交中击败了23.93%的用户