【剑指】41.数据流中的中位数

题目描述

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

算法分析

• 为了保证插入新数据和取中位数的时间效率都高效，这里使用大顶堆+小顶堆的容器，并且满足：两个堆中的数据数目差不能超过1，这样可以使中位数只会出现在两个堆的交接处；大顶堆的所有数据都小于小顶堆，这样就满足了排序要求。

提交代码(模拟最大最小堆,其实是实现了priority_queue)：

class Solution {
public:
	void Insert(int num)
	{
		int length = maxHeap.size() + minHeap.size();
		// 奇数个插入max
		if ((length & 1) == 1)
		{
			// 如果插入最大堆中的数比最小堆的最小值要大
			if (!minHeap.empty() && num > minHeap[0])
			{
				minHeap.push_back(num);
				push_heap(minHeap.begin(), minHeap.end(), greater<int>());

				num = minHeap[0];
				pop_heap(minHeap.begin(), minHeap.end(), greater<int>());
				minHeap.pop_back();
			}
			maxHeap.push_back(num);
			push_heap(maxHeap.begin(), maxHeap.end(), less<int>());
		}
		else
		{
			if (!maxHeap.empty() && num < maxHeap[0])
			{
				maxHeap.push_back(num);
				push_heap(maxHeap.begin(), maxHeap.end(), less<int>());

				num = maxHeap[0];
				pop_heap(maxHeap.begin(), maxHeap.end(), less<int>());
				maxHeap.pop_back();
			}

			minHeap.push_back(num);
			push_heap(minHeap.begin(), minHeap.end(), greater<int>());
		}
	}

	double GetMedian()
	{
		int length = maxHeap.size() + minHeap.size();
		if (length == 0)
			return 0;
		return (length & 1) == 1 ? minHeap[0] 
			: (maxHeap[0] + minHeap[0]) / 2.0;
	}

private:
	// 奇数时插入
	vector<int> maxHeap;
	// 偶数时插入
	vector<int> minHeap;
};

测试代码：

// ====================测试代码====================
void Test(char* testName, Solution numbers, double expected)
{
	if (testName != nullptr)
		printf("%s begins: ", testName);

	if (abs(numbers.GetMedian() - expected) < 0.0000001)
		printf("Passed.\n");
	else
		printf("FAILED.\n");
}

int main(int argc, char* argv[])
{
	Solution numbers;

	printf("Test1 begins: ");
	try
	{
		numbers.GetMedian();
		printf("Passed.\n");
	}
	catch (const exception&)
	{
		printf("FAILED.\n");
	}

	numbers.Insert(5);
	Test("Test2", numbers, 5);

	numbers.Insert(2);
	Test("Test3", numbers, 3.5);

	numbers.Insert(3);
	Test("Test4", numbers, 3);

	numbers.Insert(4);
	Test("Test6", numbers, 3.5);

	numbers.Insert(1);
	Test("Test5", numbers, 3);

	numbers.Insert(6);
	Test("Test7", numbers, 3.5);

	numbers.Insert(7);
	Test("Test8", numbers, 4);

	numbers.Insert(0);
	Test("Test9", numbers, 3.5);

	numbers.Insert(8);
	Test("Test10", numbers, 4);

	return 0;
}