PriorityQueue实现大小顶堆 解决topK问题

发表于 | 分类于 | 本文总阅读量 次

java.lang.Object
——继承者 java.util.AbstractCollection
————继承者 java.util.AbstractQueue
——————继承者 java.util.PriorityQueue

类型参数： E - collection 中所保存元素的类型。

所有已实现的接口： Serializable, Iterable, Collection, Queue

1
2
3
4
public class PriorityQueue<E> extends AbstractQueue<E> 
    implements java.io.Serializable {
	...
}

Demo

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
public class PriorityQueueDemo {

    public static void main(String[] args) {
        Integer[] arr = {3, 2, 5, 8, 71, 3, 69, 5, 4, 9, 82, 5, 55, 96, 485, 357};
        System.out.println(topKMax(Arrays.asList(arr), arr.length));
        System.out.println(topKMin(Arrays.asList(arr), arr.length));
    }

    public static <T extends Comparable<T>> List<T> topKMax(List<T> list, int top) {
        if (list == null) {
            throw new RuntimeException("list is null!");
        }
        PriorityQueue<T> queue = new PriorityQueue<>(top);
        for (T t : list) {
            if (queue.size() < top) {
                queue.add(t);
            } else if (queue.peek().compareTo(t) < 0) {
                queue.poll();
                queue.add(t);
            }
        }
        LinkedList<T> arr = new LinkedList<>();
        while (!queue.isEmpty()) {
            arr.addFirst(queue.poll());
        }
        return arr;
    }

    public static <T extends Comparable<T>> List<T> topKMin(List<T> list, int top) {
        if (list == null) {
            throw new RuntimeException("list is null!");
        }
        PriorityQueue<T> queue = new PriorityQueue<>(top, new Comparator<T>() {
            @Override
            public int compare(T o1, T o2) {
                return o2.compareTo(o1);
            }
        });
        for (T t : list) {
            if (queue.size() < top) {
                queue.add(t);
            } else if (queue.peek().compareTo(t) > 0) {
                queue.poll();
                queue.add(t);
            }
        }
        LinkedList<T> arr = new LinkedList<>();
        while (!queue.isEmpty()) {
            arr.addFirst(queue.poll());
        }
        return arr;
    }
}
0%