堆排序、topK

randian666 · randian666 · commit 6939177ea15f · 2019-01-31T13:47:55.000+08:00
diff --git a/src/main/java/com/algorithm/study/demo/datastructure/tree/Heap.java b/src/main/java/com/algorithm/study/demo/datastructure/tree/Heap.java
@@ -0,0 +1,106 @@
+package com.algorithm.study.demo.datastructure.tree;
+
+import com.alibaba.fastjson.JSON;
+
+/**
+ * 堆
+ * 堆就是一颗完全二叉树
+ * 堆中每一个节点的值都必须大于等于（或者小于等于）其子树（左右子节点）中每个节点的值。
+ * 堆顶存储的就是最大值或者最小值
+ * 用数组存储非常节省存储空间，下标为id的节点的左子节点就是i*2的节点，右子节点的下标为i*2+1。父节点就是i/2的节点。
+ *
+ * @Author: liuxun
+ * @CreateDate: 2019/1/31 上午10:30
+ * @Version: 1.0
+ */
+public class Heap {
+    private int[] data;
+    private int count;//堆中元素的个数
+    private int capacity;//堆容量
+
+    /**
+     * 初始化堆
+     * @param capacity
+     */
+    public Heap(int capacity){
+        this.capacity=capacity;
+        data=new int[capacity];
+        count=0;
+    }
+
+    /**
+     * 往堆中插入数据
+     * 从下往上的方法堆化
+     * @param d
+     */
+    public void insert(int d){
+        if (count>=capacity) return;//堆满了
+        count++;//从下标为1开始填充数据
+        data[count]=d;
+        int i=count;
+        while (i/2>0 && data[i]>data[i/2]){
+            //交换数据
+            swap(data,i,i/2);
+            i=i/2;//切到父节点
+        }
+    }
+
+    private static void swap(int[] a,int i,int n) {
+        int temp=a[i];
+        a[i]=a[n];
+        a[n]=temp;
+    }
+
+    /**
+     * 删除堆顶元素
+     * 把最后一个元素放到对顶，然后从上到下的方法堆化。
+     */
+    public void removeMax(){
+        data[1]=data[count];
+        count--;
+        int i=1;
+        heapify(data,count,i);
+    }
+
+    private static void buildHeap(int[] a, int n) {
+        for (int i = n/2; i >= 1; --i) {
+            heapify(a, n, i);
+        }
+        System.out.println(JSON.toJSONString(a));
+    }
+
+    private static void heapify(int[] a, int n, int i) {
+        while (true) {
+            int maxPos = i;
+            if (i*2 <= n && a[i] < a[i*2]) maxPos = i*2;
+            if (i*2+1 <= n && a[maxPos] < a[i*2+1]) maxPos = i*2+1;
+            if (maxPos == i) break;
+            swap(a, i, maxPos);
+            i = maxPos;
+        }
+    }
+
+    // n 表示数据的个数，数组 a 中的数据从下标 1 到 n 的位置。
+    public static void sort(int[] a, int n) {
+        buildHeap(a, n);//构建堆
+        int k = n;
+        while (k > 1) {
+            swap(a, 1, k);//堆顶跟最后一个元素交换
+            --k;//每次k-1，也就去取出
+            heapify(a, k, 1);//1-k堆化
+        }
+        System.out.println(JSON.toJSONString(a));
+    }
+
+
+    public static void main(String[] args) {
+//        Heap heap=new Heap(6);
+//        for (int i=1;i<=5;i++){
+//            heap.insert(i);
+//        }
+//        System.out.println(heap.count);
+
+        int[] d = new int[]{0, 5, 2, 1, 4, 3};
+        sort(d,d.length-1);
+    }
+}
diff --git a/src/main/java/com/algorithm/study/demo/datastructure/tree/TopkCount.java b/src/main/java/com/algorithm/study/demo/datastructure/tree/TopkCount.java
@@ -0,0 +1,49 @@
+package com.algorithm.study.demo.datastructure.tree;
+
+import com.alibaba.fastjson.JSON;
+
+import java.util.PriorityQueue;
+
+/**
+ * @Author: liuxun
+ * @CreateDate: 2019/1/30 下午8:10
+ * @Version: 1.0
+ */
+public class TopkCount {
+    /**
+     * 求数据中前K大数据
+     *
+     * @param data
+     * @param k
+     * @return
+     */
+    public static int[] topk(int[] data, int k) {
+        PriorityQueue<Integer> queue = new PriorityQueue<>(k);
+
+        for (int i = 0; i < data.length; i++) {
+            if (queue.size() < k) {
+                queue.offer(data[i]);
+            } else {
+                int value = queue.peek();
+                // 如果发现数据比堆顶元素大，则加入到小顶堆中
+                if (data[i] > value) {
+                    queue.poll();//删除
+                    queue.offer(data[i]);
+                }
+            }
+        }
+
+        int[] result = new int[k];
+        int index = 0;
+        // 遍历完成后，小顶堆的数据就为需要求得的topk的数据
+        while (!queue.isEmpty()) {
+            result[index++] = queue.poll();
+        }
+
+        return result;
+    }
+
+    public static void main(String[] args) {
+        System.out.println(JSON.toJSONString(topk(new int[]{1,2,3,4,5},2)));
+    }
+}
