使用PriorityQueue实现topK

package com.github.ralgond.yousuggest;

import java.util.Comparator;
import java.util.PriorityQueue;

public class Test0 {

	public static void main(String args[]) {
		int topK = 5;
		
		Comparator<Integer> cmp = new Comparator<Integer>() {

			@Override
			public int compare(Integer o1, Integer o2) {
				if (o2 > o1)
					return -1;
				else if (o2 < o1)
					return 1;
				return 0;
			}
		};
		
		PriorityQueue<Integer> pq = new PriorityQueue<>(topK, cmp);
		
		
		for (int i = 0; i < 10; i++) {
			if (pq.size() < topK)
				pq.add(i);
			else {
				if (pq.comparator().compare(i, pq.peek()) > 0) {
					pq.poll();
					pq.add(i);
				}
			}
		}
		
		for (int i : pq) {
			System.out.println(i);
		}
	}
}

posted @ 2020-12-30 12:01  ralgo  阅读(129)  评论(0)    收藏  举报