80 随时找到数据流的中位数 - 程序员代码面试指南

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.util.PriorityQueue;
import java.util.Scanner;


/**
 * 基于两个优先队列(底层是堆)实现中位数O(1)输出
 */
class MedianHolder{

    PriorityQueue<Integer> minQue;
    PriorityQueue<Integer> maxQue;
    Double median;

    public MedianHolder() {
        // 大根堆，存储小于 median 的
        minQue = new PriorityQueue<>((a,b)->{
            if(a > b){
                return -1;
            }
            if(a < b){
                return 1;
            }
            return 0;
        });
        // 小根堆，存储大于 median 的
        maxQue = new PriorityQueue<>((a,b)->{
            if(a < b){
                return -1;
            }
            if(a > b){
                return 1;
            }
            return 0;
        });
        // 初始无 median
        median = null;
    }

    public void putVal(int val){
        if(minQue.isEmpty()){
            minQue.add(val);
            median = 1.0 * val;
        }else {
            int minTopVal = minQue.peek();
            if(val <= minTopVal){
                minQue.offer(val);
            }else {
                maxQue.offer(val);
            }
            int minSize = minQue.size();
            int maxSize = maxQue.size();
            if(minSize - maxSize > 1){
                int tmpMinTop = minQue.poll();
                maxQue.offer(tmpMinTop);
                minSize --;
                maxSize ++;
            }else{
                if(maxSize - minSize > 1){
                    int tmpMaxTop = maxQue.poll();
                    minQue.offer(tmpMaxTop);
                    minSize ++;
                    maxSize --;
                }
            }
            // 重新设置 median
            if(minSize < maxSize){
                median = 1.0 * maxQue.peek();
            }else if(minSize > maxSize){
                median = 1.0 * minQue.peek();
            }else {
                median = 0.5 * (minQue.peek() + maxQue.peek());
            }
        }
    }

    public Double getMedian(){
        return median;
    }
}

public class Main {
    public static Scanner sc = new Scanner(System.in);

    public static BufferedReader reader = new BufferedReader(new InputStreamReader(System.in));

    public static void main(String[] args) throws Exception {
        int n = sc.nextInt();
        MedianHolder medianHolder = new MedianHolder();
        for(int i=0;i<n;i++){
            int op = sc.nextInt();
            if(op == 1){
                int val = sc.nextInt();
                medianHolder.putVal(val);
            }else {
                Double median = medianHolder.getMedian();
                if(median == null){
                    System.out.println(-1);
                }else {
                    System.out.println(String.format("%.1f", median));
                }
            }
        }

    }

}
/*
8
1 5
2
1 3
2
1 6
2
1 7
2

5.0
4.0
5.0
5.5
 */
ac
