最小堆解決Top K問題
阿新 • • 發佈:2018-12-30
問題描述:
有一組資料n個,要求取出這組資料中最大的K個值。
對於這個問題,解法有很多中。比如排序及部分排序,不過效率最高的要數最小堆,它的時間複雜度為O(nlogk)。
解題思路:
- 取出陣列的前n個元素,建立長度為n的最小堆。
- 從n開始迴圈陣列的剩餘元素,如果元素(a)比最小堆的根節點大,將a設定成最小堆的根節點,並讓堆保持最小堆的特性。
- 迴圈完成後,最小堆中的所有元素就是需要找的最大的n個元素。
Java實現:
1.建堆程式碼:
public class Heap<T> {
//以陣列形式儲存堆元素
private T[] heap;
//用於比較堆中的元素。c.compare(根,葉子) > 0
//使用相反的Comparator可以建立最大堆、最小堆
private Comparator<? super T> comparator;
//初始化堆
public Heap(T[] heap, Comparator<? super T> comparator) {
this.heap = heap;
this.comparator = comparator;
buildHeap();
}
//返回值為i/2
private int parent(int i) {
return (i - 1) >> 1;
}
//返回指定節點的left子節點陣列索引。相當於2*(i+1)-1
private int left(int i) {
return ((i + 1) << 1) - 1;
}
//返回指定節點的right子節點陣列索引。相當於2*(i+1)
private int right(int i) {
return (i + 1) << 1;
}
//堆化,i是堆化的起始節點
private void heapify(int i) {
heapify(i, heap.length);
}
//堆化,i是堆化的起始節點,size是堆化的範圍
private void heapify(int i, int size) {
int l = left(i);
int r = right(i);
int next = i;
if (l < size && comparator.compare(heap[l], heap[i]) > 0)
next = l;
if (r < size && comparator.compare(heap[r], heap[next]) > 0)
next = r;
if (i == next)
return;
swap(i, next);
heapify(next, size);
}
//對堆排序
public void sort() {
// buildHeap();
for (int i = heap.length - 1; i > 0; i--) {
swap(0, i);
heapify(0, i);
}
}
//交換陣列值
private void swap(int i, int j) {
T tmp = heap[i];
heap[i] = heap[j];
heap[j] = tmp;
}
//建立堆
private void buildHeap() {
for (int i = (heap.length) / 2 - 1; i >= 0; i--) {
heapify(i);
}
}
public void setRoot(T root) {
heap[0] = root;
heapify(0);
}
public T root() {
return heap[0];
}
}
2.最小堆實現Top K篩選
public class TopK {
public static int[] toPrimitive(Integer array[]) {
if (array == null)
return null;
if (array.length == 0)
return new int[0];
int result[] = new int[array.length];
for (int i = 0; i < array.length; i++){
result[i] = array[i].intValue();
}
return result;
}
public static int[] topK(int[] array, int n) {
if (n >= array.length) {
return array;
}
Integer[] topn = new Integer[n];
for (int i = 0; i < topn.length; i++) {
topn[i] = array[i];
}
Heap<Integer> heap = new Heap<Integer>(topn, new Comparator<Integer>() {
@Override
public int compare(Integer o1, Integer o2) {
// 生成最大堆使用o1-o2,生成最小堆使用o2-o1
return o2 - o1;
}
});
for (int i = n; i < array.length; i++) {
int value = array[i];
int min = heap.root();
if (value > min) {
heap.setRoot(value);
}
}
return toPrimitive(topn);
}
public static void main(String[] args) {
Random random = new Random(100);
int[] array = new int[100];
for(int i=0; i<100; i++){
array[i] = random.nextInt(1000);
}
int count = 0;
for(int i=0; i<100; i++){
count++;
System.out.print(array[i] + " ");
if(count >= 25){
System.out.println();
count=0;
}
}
System.out.print("Top 10 :");
int[] result = topK(array, 10);
for (int integer : result) {
System.out.print(integer + ",");
}
}
}