剑指 Offer 41. 数据流中的中位数

题目描述

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

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

设计一个支持以下两种操作的数据结构：

• void addNum(int num) - 从数据流中添加一个整数到数据结构中。
• double findMedian() - 返回目前所有元素的中位数。

示例 1：

输入：
["MedianFinder","addNum","addNum","findMedian","addNum","findMedian"]
[[],[1],[2],[],[3],[]]
输出：[null,null,null,1.50000,null,2.00000]

示例 2：

输入：
["MedianFinder","addNum","findMedian","addNum","findMedian"]
[[],[2],[],[3],[]]
输出：[null,null,2.00000,null,2.50000]

限制：

• 最多会对 addNum、findMedia进行 50000 次调用。

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/shu-ju-liu-zhong-de-zhong-wei-shu-lcof

代码实现

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

	void addNum(int num) {
		if (largeRootHeap.size() <= smallRootHeap.size()) {
			insertHeap(smallRootHeap, num, -1);
			int midValue = popHeapTop(smallRootHeap, -1);
			insertHeap(largeRootHeap, midValue, 1);
		}
		else {
			insertHeap(largeRootHeap, num, 1);
			int midValue = popHeapTop(largeRootHeap, 1);
			insertHeap(smallRootHeap, midValue, -1);
		}
		return;
	}

	double findMedian() {
		double result;
		if (largeRootHeap.size() == smallRootHeap.size())
			result = (1.0 * (largeRootHeap[0] + smallRootHeap[0])) / 2;
		else
			result = largeRootHeap[0];
		return result;
	}

private:
	vector<int> largeRootHeap;
	vector<int> smallRootHeap;

	void insertHeap(vector<int>& heap, int value, int type) {
		// type = 1: Large Heap
		// type = -1: Small Heap
		heap.push_back(value);
		int valueIndex = heap.size() - 1;
		int parentIndex = ((valueIndex + 1) / 2) - 1;
		while (true) {
			if (parentIndex < 0 || (value - heap[parentIndex]) * type < 0)
				break;
			swap(heap[valueIndex], heap[parentIndex]);
			valueIndex = parentIndex;
			parentIndex = ((valueIndex + 1) / 2) - 1;
		}
		return;
	}

	int popHeapTop(vector<int>& heap, int type) {
		// type = 1: Large Heap
		// type = -1: Small Heap
		int top = heap[0];
		heap[0] = heap[heap.size() - 1];
		int value = heap[0];
		heap.pop_back();
		int valueIndex = 0;
		int childIndex1 = 2 * valueIndex + 1;
		int childIndex2 = 2 * valueIndex + 2;
		int childIndex;
		if (childIndex1 < heap.size() && childIndex2 < heap.size())
			childIndex = ((heap[childIndex1] - heap[childIndex2]) * type > 0) ? childIndex1 : childIndex2;
		else
			childIndex = childIndex1;
		while (true) {
			if (childIndex >= heap.size() || (value - heap[childIndex]) * type > 0)
				break;
			swap(heap[valueIndex], heap[childIndex]);
			valueIndex = childIndex;
			childIndex1 = 2 * valueIndex + 1;
			childIndex2 = 2 * valueIndex + 2;
			if (childIndex1 < heap.size() && childIndex2 < heap.size())
				childIndex = ((heap[childIndex1] - heap[childIndex2]) * type > 0) ? childIndex1 : childIndex2;
			else
				childIndex = childIndex1;
		}
		return top;
	}
};

/**
 * Your MedianFinder object will be instantiated and called as such:
 * MedianFinder* obj = new MedianFinder();
 * obj->addNum(num);
 * double param_2 = obj->findMedian();
 */

解题思路

1. 使用一个大根堆largeRootHeap存储数组的较小一半，使用一个小根堆smallRootHeap存储数组的较大一半；
2. 判断大根堆与小根堆的元素个数，若大根堆元素个数较少或两堆元素个数相等，则希望大根堆的元素+1，否则希望小根堆的元素+1；
3. 大根堆的元素+1：向小根堆插入新元素，并取出小根堆堆顶元素插入大根堆；
4. 小根堆的元素+1：向大根堆插入新元素，并取出大根堆堆顶元素插入小根堆；
5. 取出中位数：若两堆元素相等，取两堆顶元素平均值，否则取大根堆堆顶。