1. 程式人生 > >Leetcode 295. Find Median from Data Stream

Leetcode 295. Find Median from Data Stream

用了兩個優先佇列

最後結果比 91% java Submission 快。

一開始的思路,如果要求速度儘量快的話,肯定不能插入和查詢分開來做 ,要在插入的時候就對這堆東西進行排序,並找出中值,所以一開始想到的是二叉平衡樹。

但是二叉平衡樹只能保證左右子樹的高度差距不大於1,並不能保證兩邊的節點個數相等。並且,在處理總結點個數為偶數的時候有困難。

所以不如將左右子樹分開來處理,分別記錄兩棵樹的大小。

然後發現其實是兩個優先佇列組合使用的問題。

左邊是一個大數先出的優先佇列,右邊是一個小數先出的優先佇列。且右邊佇列裡的所有值均不小於左邊的。

實現起來的時候發現自己已經把優先佇列忘記的差不多了。。想想怎麼把BigTree和SmallTree優化成一個數據結構。。

直接呼叫PriorityQueue的話就更方便了。。但也算複習一下吧。。

很久沒刷題了果然變得很菜很菜。。

class MedianFinder {

    class BigTree {  
        public int size = 0;
        private int[] array = new int[10000];
        public int rootVal = 0;
        public void add(int val) {
        	array[++size] = val;
        	swim(size);
            rootVal = array[1];
        }
        public int remove() {
        	int val = array[1];
        	exch(1, size--);
        	array[size+1] = 0;
        	sink(1);
        	
            rootVal = array[1];
            
            return val;
        }
        private void exch(int i, int j) {
        	int temp = array[i];
        	array[i] = array[j];
        	array[j] = temp;
        }
        private void swim(int k) {
        	while(k > 1 && array[k/2] > array[k]) {
        		exch(k/2, k);
        		k = k/2;
        	}
        }
        private void sink(int k) {
        	while(2*k <= size) {
        		int j = 2*k;
        		if(j < size && array[j] > array[j+1]) j++;
        		if(array[k] <= array[j]) break;
        		exch(k, j);
        		k = j;
        	}
        }  
    }
    
    class SmallTree {
    	public int size = 0;
        private int[] array = new int[10000];
        public int rootVal = 0;
        public void add(int val) {
        	array[++size] = val;
        	swim(size);
            rootVal = array[1];
        }
        public int remove() {
        	int val = array[1];
        	exch(1, size--);
        	array[size+1] = 0;
        	sink(1);
        	
            rootVal = array[1];
            return val;
        }
        private void exch(int i, int j) {
        	int temp = array[i];
        	array[i] = array[j];
        	array[j] = temp;
        }
        private void swim(int k) {
        	while(k > 1 && array[k/2] < array[k]) {
        		exch(k/2, k);
        		k = k/2;
        	}
        }
        private void sink(int k) {
        	while(2*k <= size) {
        		int j = 2*k;
        		if(j < size && array[j] < array[j+1]) j++;
        		if(array[k] >= array[j]) break;
        		exch(k, j);
        		k = j;
        	}
        }  
    }
    
    private BigTree bigTree;
    private SmallTree smallTree;
    
    /** initialize your data structure here. */
    public MedianFinder() {
        bigTree = new BigTree();
        smallTree = new SmallTree();
    }
    
    public void addNum(int num) {
        if(num > bigTree.rootVal) bigTree.add(num);
        else if(num <= bigTree.rootVal) smallTree.add(num);

        if(bigTree.size - smallTree.size > 1) {
            int temp = bigTree.remove();
            smallTree.add(temp);
        }
        else if(smallTree.size - bigTree.size > 1) {
            int temp = smallTree.remove();
            bigTree.add(temp);
        }
    }
    
    public double findMedian() {
        if(bigTree.size == smallTree.size) return ((double)(bigTree.rootVal + smallTree.rootVal))/2;
        else if(bigTree.size > smallTree.size) return bigTree.rootVal;
        else return smallTree.rootVal;
    }
}

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