用PriorityQueue解決選擇最小的K個數問題

小樓昨夜又西風發表於2019-03-18

在這裡插入圖片描述
Java沒有堆的實現,現場寫也來不及,有的文獻說用TreeSet,比如劍指offer,但是TreeSet是一個set,相同的數只能存一個,個人感覺不合適,相比之下,Java中的PriorityQueue倒是一個不錯的選擇。用PriorityQueue的實現的程式碼如下:

/**
	 * 用PriorityQueue實現選擇最小的k個數
	 * @param array	陣列
	 * @param k
	 * @return
	 */
	public int[] selectKmin(int[] array,int k){
		int[] res = new int[k];	
		//建立一個降序排列的PriorityQueue,自定義比較器作為引數
		PriorityQueue<Integer> priorityQueue = new PriorityQueue<Integer>(k,new Comparator<Integer>(){
 
			@Override
			public int compare(Integer o1, Integer o2) {
				return o2-o1;
			}});
		
		for(int i=0;i<array.length;i++){
			if(priorityQueue.size()<k){
				priorityQueue.add(array[i]);
			}else{
				int maxInQueue = priorityQueue.peek();
				if(maxInQueue>array[i]){//每次取數與堆頂的元素進行比較,
				//如果堆頂元素大,則刪除堆頂元素,並新增這個新數到堆中
					priorityQueue.poll();
					priorityQueue.add(array[i]);
				}
			}
		}
		
		//升序存放
		for(int i=0;i<k;i++){
			res[k-i-1] = priorityQueue.poll();
		}
		
		return res;
	}


相關文章