使用堆来查找数据流中的中位数

思想:维护一个最大堆和最小堆,最大堆来维护小于中位数的数字,最小堆用来维护大于中位数的数字,在遍历完所有数据后,两个堆中的顶端的元素肯定在整个数据流中最中间的位置,由此可得,思想跟快排中维护less和more思想是相似的,但是还需要注意维护两个边界条件:

  • 最大堆的最大值应小于最小堆的最小值
  • 最大堆和最小堆的“重量差”——即两个堆中的数字个数的差不能超过1,这个也是区别于快排,能够找到中位数的关键
import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;

/**
 * 
 * 使用堆来找数据流中的中位数
 */
public class FindMediumByHeap {
	private static PriorityQueue maxHeap;	// 最大堆存放比中位数小的数
	private static PriorityQueue minHeap;	// 最小堆存放比中位数大的数
	
	private static int[] generateArray(int N) {
		int[] arr = new int[N];
		for (int i = 0; i < arr.length; i++) {
			arr[i] = (int) (Math.random() * (N+1));
		}
		return arr;
	}
	
	private static void printArray(int[] arr) {

		if (arr == null) {
			return;
		}
		for(int i=0;i();	
		maxHeap = new PriorityQueue<>(arr.length, new ComparatorBigHeap());
		
		for (int n : arr) {
			
			// 当堆为空时先插入堆
			if(maxHeap.isEmpty()) {
				maxHeap.add(n);
				continue;
			} else if(minHeap.isEmpty()) {
				minHeap.add(n);
				continue;
			}
			
			
			// 调整 less 和 more 的边界, 处理less大于more的情况,交换一下边界好了
			if(maxHeap.peek() > minHeap.peek()) {
				int tmp = minHeap.remove();
				minHeap.add(maxHeap.remove());
				maxHeap.add(tmp);
			}
			
			if(n <= maxHeap.peek()) {
				maxHeap.add(n);
			} else {
				minHeap.add(n);
			}
			
			// 调整两堆的平衡
			if(maxHeap.size()-minHeap.size() >= 2) {
				minHeap.add(maxHeap.remove());
			} else if (minHeap.size()-maxHeap.size() >= 2) {
				maxHeap.add(minHeap.remove());
			}
		}
		
		// 存放完毕
		if(maxHeap.size() > minHeap.size()) { // 小于中位数的多
			return maxHeap.remove();
		} else if(minHeap.size() > maxHeap.size()) { // 大于等于中位数的多
			return minHeap.remove();
		} else {
			return ((float)minHeap.remove() + (float)maxHeap.remove())/2;
		}
	}
	
	public static void main(String[] args) {
		int N = 21;	// 参与测试的队列长度
		int[] arr = generateArray(N);
//		int[] arr = new int[]{3, 1, 20, 16, 7, 18, 14, 9, 15, 9, 2, 17, 7, 1, 9, 10, 14, 5, 8, 5 };
		printArray(arr);
		
		System.out.println(findMedium(arr));

		System.out.println("正确答案,在排序后的数组中查找");
		Arrays.sort(arr);
		printArray(arr);
		
		if(N%2 == 0) 
			System.out.println(((float)arr[N/2] + (float)arr[N/2-1])/2);
		else {
			System.out.println((float)arr[N/2]);
		}
	}
}


class ComparatorBigHeap implements Comparator {

	@Override
	public int compare(Integer o1, Integer o2) {
		// TODO Auto-generated method stub
		return o2.compareTo(o1);
	}	
}

你可能感兴趣的:(使用堆来查找数据流中的中位数)