leetcode-295

大顶堆-小顶堆.png

    static class MedianFinder {
        /**
         * 大顶堆
         *  6
         * 3 5
         */
        private PriorityQueue<Integer> bigPriorityQueue;

        /**
         * 小顶堆
         *   9
         * 10 14
         */
        private PriorityQueue<Integer> smallPriorityQueue;

        /** initialize your data structure here. */
        public MedianFinder() {
            /**
             * 小数在顶堆（小顶堆）
             */
            smallPriorityQueue = new PriorityQueue();

            /**
             * 大数在顶堆（大顶堆）
             */
            bigPriorityQueue = new PriorityQueue<>(new Comparator<Integer>() {
                @Override
                public int compare(Integer o1, Integer o2) {
                    return -(o1 - o2);
                }
            });
        }

        public void addNum(int num) {
            int bigQueueSize = bigPriorityQueue.size();
            int smallQueueSize = smallPriorityQueue.size();
            if (bigQueueSize == 0) {
                bigPriorityQueue.add(num);
                return;
            }

            if (bigQueueSize == smallQueueSize) {
                if (num > smallPriorityQueue.peek()) {
                    smallPriorityQueue.add(num);
                } else {
                    bigPriorityQueue.add(num);
                }
            } else if (bigQueueSize > smallQueueSize) {
                /**
                 * 最小堆
                 * 小数在上
                 */
                if (num > bigPriorityQueue.peek()) {
                    smallPriorityQueue.add(num);
                } else {
                    smallPriorityQueue.add(bigPriorityQueue.peek());
                    bigPriorityQueue.poll();
                    bigPriorityQueue.add(num);
                }

            } else if (bigQueueSize < smallQueueSize) {
                if (num < smallPriorityQueue.peek()) {
                    bigPriorityQueue.add(num);
                } else {
                    bigPriorityQueue.add(smallPriorityQueue.peek());
                    smallPriorityQueue.poll();
                    smallPriorityQueue.add(num);
                }
            }
        }

        public double findMedian() {
            if (smallPriorityQueue.size() == bigPriorityQueue.size()) {
                return (smallPriorityQueue.peek() + bigPriorityQueue.peek()) / 2.0D;
            } else if (smallPriorityQueue.size() > bigPriorityQueue.size()) {
                return smallPriorityQueue.peek();
            } else {
                return bigPriorityQueue.peek();
            }
        }
    }