简单地用Java解决topN问题

个人blog原文地址：http://www.gemoji.me/java_top_n/

距离上次写博客有一个月了， 反省下。今天先写篇简单点的，算是热热身吧。

写在前面
我想几乎每个找过工作的程序员都曾经在面试的时候遇到过如何求topN的问题，而且多数都能不假思索的回答：求topN大用小顶堆，求topN小用大顶堆（觉得反了的同学请去面壁。。。），但是应该也有一部分同学和我之前一样，一直只是把它作为一道面试题而已。

思路
说来不怕丢人，我真的是最近才遇到需要求top N大的场景，囧。。。当时的第一反应是要自己实现一个固定大小的二叉堆，因为印象中常用的Java工具类里并没有现成的实现，但是从头开始实现一个基础工具类是件挺麻烦的事情，别的不说，单是测试就得耗费不少功夫，网上虽然有很多现有实现，但是都不见得那么靠谱。于是想了一个折中的办法就是，对现有的成熟的Java工具类做封装，以达到想要的功能。

代码实现
JDK里有一个现成的工具类叫PriorityQueue，顾名思义，实现的是优先队列的功能，它能够保证每次插入新元素后，队列首位始终是优先级最高的元素，这正是我们想要的，但是它并没有对队列的容量做限制，因此只要对这一点做些许改造，就能达到我们期望的效果。代码如下：

	//固定容量的优先队列，模拟大顶堆，用于解决求topN小的问题
	public static class FixSizedPriorityQueue>{
		private PriorityQueue queue;
		private int  maxSize; //堆的最大容量

		public FixSizedPriorityQueue(int maxSize){
			if(maxSize <= 0) throw new IllegalArgumentException();			
			this.maxSize = maxSize;
			this.queue = new PriorityQueue(maxSize, new Comparator() {

				@Override
				public int compare(E o1, E o2) {
					return o2.compareTo(o1);
				}
			});
		}

		public void add(E e){
			if(queue.size() < maxSize){ //未达到最大容量，直接添加
				queue.add(e);
			}else{ //队列已满
				E peek = queue.peek();
				if(e.compareTo(peek) < 0){ //将新元素与当前堆顶元素比较，保留较小的元素
					queue.poll();
					queue.add(e);
				}
			}
		}

		public List sortedList(){
			List list = new ArrayList(queue);
			Collections.sort(list); //PriorityQueue本身的遍历是无序的，最终需要对队列中的元素进行排序
			return list;
		}

	}

注意一下构造PriorityQueue时用到的Comparator，这里面有个比较tricky的地方，因为代码本身是想实现一个大顶堆来解决求topN小的问题，而PriorityQueue的优先级定义是元素越小优先级越高，这与我们的期望相反，因此我们需要自定义PriorityQueue的Comparator，反转元素大小的定义，这样就能保证队列首位的元素是当前队列中最大的。而在比较新元素和队列首位元素的大小时，则是按照正常的元素大小定义做比较。

结尾
在需要快速实现需求的时候，我们可能没有足够的时间把代码写得很完美，如何写最少的代码实现想要的功能，也是我们日常工作中需要考虑的问题。
PS：文中若有纰漏之处，敬请拍砖。

评论

2 楼 lc87624 2013-08-12  

luxury_zh 写道

最后一点不明白，为什么每次还要对队列进行排序，循环N次，调用poll方法每次取得都是当前队列最大的数字了。

你说的方法也是可以的，我当时应该是想保留queue里的数据，所以做了一份copy。

1 楼 luxury_zh 2013-08-04  

最后一点不明白，为什么每次还要对队列进行排序，循环N次，调用poll方法每次取得都是当前队列最大的数字了。