原创：搜索排序算法之自定义性能优良的PriorityQueue（与Python的heap比较）

　　前几天写了一篇关于"史上对BM25模型最全面最深刻解读以及lucene排序深入解读"的博客，lucene最后排序用到的思想是"从海量数据中寻找topK"的时间空间最优算法(这是一个博士的学术论文)。在特定的场合，比如solr自带的搜索智能提示公能，当构建完三叉树，前缀匹配查找出所有的节点之后，也要用这种思想进行排序。根据这个思想构造出一个优先级队列，具有容量限制(K),精确的时间复杂度为KlgK+(n-k)lgK，最坏的时间复杂度：(n-k)*lgk +lg(k-1)!。远远优于目前任何的排序算法，在学术论文里已经进行了理论论证。今天根据这个思想，尝试着写了这个数据结构，经过修改，已经接近完美。这个数据结构的适用场景就是从海量数据中寻找出topK的数据来，可以解决lucene排序，解决搜索推荐的冷启动问题(从用户的搜索日志中寻找出topK,推荐出来)。凡是有容量限制的这类问题，都可以使用它。而TreeSet或者TreeMap的排序的时间复杂度是O(lgn)，并且底层基于可排序的二叉树(通过红黑树获取平衡)，然后中序遍历得到最后结果。在IK分词的消除歧义分词中用到了TreeSet。这两个数据结构常常用于Key-value数据形式并且容量不是很大或者没有容量限制的场景。其他的基础排序算法比如快排，堆排，mergeSort还有近几年新的排序TimSort，常常用于Array的排序。从底层深刻理解这些基础算法很重要！机器学习排序在搜索中的应用，最后仍然要用今天要写的PriorityQueue。JDK中有自带的PriorityQueue,但是没有容量限制，性能比较差。上传本人写的代码：

package com.txq.test;
/**
 * 优先级队列，从海量数据中寻找topK的时间空间最优算法，时间复杂度为n * lgK,空间为K,其中K << n.
 * @author XueQiang Tong
 * @since JDK1.8
 * @param <T>
 */
public abstract class PriorityQueue<T> {
	protected T heap[];//堆
	protected int size;//容量
	protected int heapSize;//最大容量

	public PriorityQueue(int capacity){
		if (capacity >= (1 << 31) - 1) {
            throw new IllegalArgumentException("max size must be <= (1 << 31) -1;got:" + capacity);
        }
		this.heapSize = capacity;
		this.size = 0;
		Object o[] = new Object[this.heapSize];
		heap = (T[]) o;
	}

	public PriorityQueue(){
		this(5);
	}

	protected abstract boolean lessThan(T t, T data);
	/**
	 * 向队列中添加元素，如果没有达到容量限制，直接添加并且构建小根堆，如果超出了容量，用大于堆顶的元素替换堆顶，然后调整堆
	 * @param data
	 */
	protected synchronized final void insertWithOverFlow(T data){
		if(data == null) return;
		if(this.size < this.heapSize){
			add(data);
		} else{
			if(this.size > 0 && lessThan(heap[0],data)){
				heap[0] = data;
				minify(0);
			}
		}
	}
	/**
	 * 按照降序依次取出topK元素，第一次取时，先构建大根堆，以后直接取堆顶元素，然后重新调整大根堆
	 * @return
	 */
	protected synchronized final T pop(){
		if(this.size == 0) return null;
		if(this.size == this.heapSize){
			for(int i = this.heapSize / 2 - 1;i >= 0;i--){
				maxnify(i);
			}
		}
		T result = heap[0];
		heap[0] = heap[this.size - 1];
		heap[this.size - 1] = null;
		this.size --;
		maxnify(0);
		return result;
	}

	public final int size(){
		return this.size;
	}

	protected final T top(){
		return this.heap[0];
	}
	/**
	 * 调整大根堆
	 * @param i
	 */
	private void maxnify(int i) {
		int left = 2 * i + 1;
		int right = 2 * i + 2;
		int max;

		if(left < this.size && lessThan(heap[i],heap[left])) max = left;
		else max = i;
		if(right < this.size && lessThan(heap[max],heap[right])) max = right;

		if (max == i || max >= this.size) return;
		swap(heap,i,max);
		maxnify(max);
	}

	private void swap(T[] heap, int i, int j) {
		T tmp;
		tmp = heap[i];
		heap[i] = heap[j];
		heap[j] = tmp;
	}
	/**
	 * 调整小根堆
	 * @param i
	 */
	private void minify(int i) {
		int left = 2 * i + 1;
		int right = 2 * i + 2;
		int min;

		if(left < this.size && lessThan(heap[left],heap[i])) min = left;
		else min = i;
		if(right < this.size && lessThan(heap[right],heap[min])) min = right;

		if (min == i || min >= this.size) return;
		swap(heap,i,min);
		minify(min);
	}
	/**
	 * 添加元素并且构建小根堆
	 * @param data
	 */
	private void add(T data) {
		this.heap[this.size++] = data;
		for(int i = this.size / 2 - 1;i >= 0;i--){
			minify(i);
		}
	}

	public synchronized final void clear() {
        for (int i = 0; i <= this.size; ++i) {
            this.heap[i] = null;
        }
        this.size = 0;
    }
}

package com.txq.test;

public class MyPriorityQueue extends PriorityQueue<Integer> {

	@Override
	protected boolean lessThan(Integer node, Integer data) {
		return node < data;
	}

	public MyPriorityQueue(int capacity){
		super(capacity);
	}
}

测试类：

package chinese.utility.utils;

import org.junit.Test;

public class PriorityQueueTest {
    PriorityQueue<Integer> q;

@Test
    public void test() {
        q = new MyPriorityQueue(3);
        q.insertWithOverFlow(99);
        q.insertWithOverFlow(78);
        q.insertWithOverFlow(109);
        q.insertWithOverFlow(123);
        q.insertWithOverFlow(23);
        q.insertWithOverFlow(45);
        q.insertWithOverFlow(56);
        Integer element = (Integer) q.pop();
        while(element != null){
            System.out.println(element);
            element = (Integer) q.pop();
        }        
    }

}

运行结果：
123

109

99

另外一个场景，比如现在有一个矩阵，没行数据都是降序排列的，维度有很多，要求找出其中matrix.length个最大值，用上面的算法就不行了，时间复杂度太高了。因为每行数据都排序好了，可以采取以下策略：

package com.txq.test;

import java.util.ArrayList;
import java.util.List;

public class PriotiryQueueTest {

	public static void main(String[] args) {
		//从以下矩阵中找出matrix.length个最大值，时间复杂度为k*lgk（k << n)
		int matrix[][] = {{100,89,78,73,69,58,40},{98,91,82,80,39,28,16},{88,80,79,76,68,60,59},
				{120,110,98,76,68,60,59},{110,100,99,67,68,60,59},{87,79,74,67,68,60,59}};
		final int heapSize = 6;
		NumInfo heap[] = new NumInfo[heapSize];
		int size = 0;
		List<Integer> result = new ArrayList<Integer>(heapSize);
		NumInfo ni;
		//1.初始化
		for(int i = 0;i < heapSize;i++){
			ni = new NumInfo(matrix[i][0],i,0);
			heap[size++] = ni;
		}
		//2.构建大根堆
		for(int j = heap.length / 2 - 1;j >= 0;j--){
			maxnify(heap,j,size);
		}
		//3.取出堆顶元素
		result.add(heap[0].data);
		heap[0].label ++;
		//4.开始迭代
		for(int i = 1;i < heapSize;i++){
			int index = heap[0].index;
			int label = heap[0].label;
			System.out.println("index:"+index+" label:"+label);
			ni = new NumInfo(matrix[index][label],index,label);
			heap[0] = ni;
			maxnify(heap,0,size);
			result.add(heap[0].data);
			heap[0].label ++;
		}
		for(int i : result){
			System.out.println(i);
		}
 	}

	/**
	 * 调整大根堆
	 * @param i
	 */
	private static void maxnify(NumInfo heap[],int i,int size) {
		int left = 2 * i + 1;
		int right = 2 * i + 2;
		int max;

		if(left < size && lessThan(heap[i],heap[left])) max = left;
		else max = i;
		if(right < size && lessThan(heap[max],heap[right])) max = right;

		if (max == i || max >= size) return;
		swap(heap,i,max);
		maxnify(heap,max,size);
	}

	private static void swap(NumInfo[] heap, int i, int max) {
		NumInfo tmp;
		tmp = heap[i];
		heap[i] = heap[max];
		heap[max] = tmp;
	}

	private static boolean lessThan(NumInfo numInfo, NumInfo numInfo2) {
		return numInfo.data < numInfo2.data;
	}

}
class NumInfo {
	public int data;
	public int index;
	public int label;

	public NumInfo(int data,int index,int label){
		this.data = data;
		this.index = index;
		this.label = label;
	}
}

另外，Python中也有类似于优先级队列的数据结构，JDK中有自带的优先级队列，都没有容量限制。Python更加倾向于函数式编程。Python的实现是靠堆操作函数的模块，叫heapq。今天试着使用一下：

from heapq import *;
from random import shuffle;
data = [x for x in range(10)];
shuffle(data);
heap = [];
for n in data:
    heappush(heap,n);

print(heap);
print(type(heappop(heap)))
print(nlargest(5,heap))#输出前5个最大值
print(nsmallest(5,heap))#输出前5个最小值
[0, 1, 6, 3, 2, 7, 9, 5, 4, 8]
<class 'int'>
[9, 8, 7, 6, 5]
[1, 2, 3, 4, 5]
可以看出，Python比Java更加灵活！从海量数据中寻找出topK问题的最优解是前面写的优先级队列解决方案，Python仍然可以完成这个功能，现在来模拟这个场景，对Python中 的heap增加容量限制：

#从海量数据中找出top4

from heapq import *;

data = [2,2,6,7,9,12,34,0,76,-12,45,79,102];#模拟海量数据
s = set();
#首先从海量数据中构造出容量为4的set，然后加载到heap中
for num in data:
    s.add(data.pop(0));
    if s.__len__() == 4:
        break;

heap = [];
for n in s:
    heappush(heap,n);

print(heap);

for num in data:
    if num > heap[0]:
        heapreplace(heap,num);#对剩余的海量数据继续迭代，如果比堆顶元素大的话，替换之并且调整小根堆

print(nlargest(4,heap))#输出前4个最大值，最后输出的时候执行堆排序！

[2, 7, 6, 9]
[102, 79, 76, 45]