leetcode Find Median from Data Stream

题目连接

https://leetcode.com/problems/find-median-from-data-stream/ 

Find Median from Data Stream

Description

Median is the middle value in an ordered integer list. If the size of the list is even, there is no middle value. So the median is the mean of the two middle value.

Examples: 
[2,3,4] , the median is 3

[2,3], the median is (2 + 3) / 2 = 2.5

Design a data structure that supports the following two operations:

void addNum(int num) - Add a integer number from the data stream to the data structure. 
double findMedian() - Return the median of all elements so far. 
For example:

add(1) 
add(2) 
findMedian() -> 1.5 
add(3) 
findMedian() -> 2

平衡树裸题。。

class MedianFinder {
public:
	const int INF = ~0u >> 1;
	MedianFinder() { null = new Node(INF, 0, NULL); root = null; }
	~MedianFinder() { clear(root), delete null; }
	void addNum(int num) {
		insert(root, num);
	}
	double findMedian() {
		int t = root->s;
		if (t & 1) {
			return (double)kth((t >> 1) + 1);
		} else {
			int v1 = kth(t >> 1);
			int v2 = kth((t >> 1) + 1);
			return (v1 + v2) / 2.0;
		}
	}
private:
	struct Node {
		int v, s, c;
		Node *ch[2];
		Node() = default;
		Node(int _v_, int _s_, Node *p) {
			v = _v_, s = c = _s_;
			ch[0] = ch[1] = p;
		}
		inline void push_up() {
			s = ch[0]->s + ch[1]->s + c;
		}
		inline int cmp(int x) const {
			return x == v ? -1 : x > v;
		}
	}*root, *null;
	inline Node *newNode(int v) {
		Node *x = new Node(v, 1, null);
		return x;
	}
	inline void rotate(Node *&x, int d) {
		Node *k = x->ch[!d]; x->ch[!d] = k->ch[d], k->ch[d] = x;
		k->s = x->s; x->push_up(); x = k;
	}
	inline void Maintain(Node *&x, int d) {
		if (!x->ch[d]->s) return;
		if (x->ch[d]->ch[d]->s > x->ch[!d]->s) rotate(x, !d);
		else if (x->ch[d]->ch[!d]->s > x->ch[!d]->s) rotate(x->ch[d], d), rotate(x, !d);
		else return;
		Maintain(x, 0), Maintain(x, 1);
	}
	inline void insert(Node *&x, int v) {
		if (!x->s) { x = newNode(v); return; }
		x->s++;
		int d = x->cmp(v);
		if (-1 == d) { x->c++; return; }
		insert(x->ch[d], v);
		x->push_up();
		Maintain(x, d);
	}
	inline void clear(Node *&x) {
		if (x->s) clear(x->ch[0]), clear(x->ch[1]), delete x;
	}
	inline int kth(int k) {
		Node *x = root;
		for (int t = 0; x->s;) {
			t = x->ch[0]->s;
			if (k <= t) x = x->ch[0];
			else if (t + 1 <= k && k <= t + x->c) break;
			else k -= t + x->c, x = x->ch[1];
		}
		return x->v;
	}
};