java 使用优先级队列管理元素中值的数据结构的Java实现（在O（1）中查看，在O（logN）中插入/删除）

Posted 2021-05-08

import java.util.*;

public class MedianManager<T extends Number & Comparable<T>> {

	private PriorityQueue<T> maxPQ;
	private PriorityQueue<T> minPQ;

	public MedianManager() {
		maxPQ = new PriorityQueue<T>((T e1, T e2) -> -e1.compareTo(e2));
		minPQ = new PriorityQueue<T>((T e1, T e2) -> e1.compareTo(e2));
	}

	public void insert(T element) {
		if (maxPQ.size() == 0) {
			maxPQ.add(element);
		}
		else if (maxPQ.size() > minPQ.size()){
			assert(maxPQ.size() - minPQ.size() == 1);
			if (!isLessThan(element, maxPQ.peek())) {
				minPQ.add(element);
			}
			else {
				maxPQ.add(element);
				minPQ.add(maxPQ.poll());
			}
			assert(maxPQ.size() == minPQ.size());
		}
		else {
			assert(maxPQ.size() == minPQ.size());
			if (!isGreaterThan(element, maxPQ.peek())) {
				maxPQ.add(element);
			}
			else {
				minPQ.add(element);
				maxPQ.add(minPQ.poll());
			}
			assert(maxPQ.size() - minPQ.size() == 1);
		}
	}

	public double peekMedian() {
		if (isEmpty()) {
			throw new RuntimeException("trying to peek while there are no elements");
		}

		if (maxPQ.size() > minPQ.size()) {
			return maxPQ.peek().doubleValue();
		}
		else {
			return (maxPQ.peek().doubleValue() + minPQ.peek().doubleValue()) / 2;
		}
	}

	public double removeMedian() {
		if (isEmpty()) {
			throw new RuntimeException("trying to remove while there are no elements");
		}

		if (maxPQ.size() > minPQ.size()) {
			return maxPQ.poll().doubleValue();
		}
		else {
			return (maxPQ.poll().doubleValue() + minPQ.poll().doubleValue()) / 2;
		}
	}

	public boolean isEmpty() {
		return maxPQ.isEmpty() && minPQ.isEmpty();
	}

	private boolean isGreaterThan(T a, T b) {
		return a.compareTo(b) > 0;
	}

	private boolean isLessThan(T a, T b) {
		return a.compareTo(b) < 0;
	}

	// sample test code below

	public static void main(String[] args) {
		MedianManager<Integer> median = new MedianManager<>();
		median.insert(1);
		System.out.println(median.peekMedian());
		median.insert(2);
		System.out.println(median.peekMedian());
		median.insert(3);
		System.out.println(median.peekMedian());
		median.insert(4);
		System.out.println(median.peekMedian());
		median.insert(5);
		System.out.println(median.removeMedian());
		System.out.println(median.removeMedian());
		System.out.println(median.removeMedian());
	}

}