堆的应用：牛客网-求最小的K个数

题目

求最小的k个数
给定一个数组，找出其中最小的K个数。例如数组元素是4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4。如果K>数组的长度，那么返回一个空的数组

输入

[4,5,1,6,2,7,3,8],4
输出
[1,2,3,4]

分析

这是一个典型的堆应用问题，可以使用大顶堆来解决这个问题

1. 拿输入数组的前K个数，来构建一个大顶堆
2. 遍历数组，更新大顶堆
3. 堆排序，输出有序数组

源码

import java.util.ArrayList;

public class Solution {
    public ArrayList<Integer> GetLeastNumbers_Solution(int[] input, int k) {
        if (k == 0 || input == null || k > input.length) {
            return new ArrayList<Integer>();
        } else {
            int[] heap = new int[k + 1];
            for (int i = 0; i < k; i++) {
                heap[i + 1] = input[i];
            }
            //初始化大顶堆
            heap(k, heap);
            //更新大顶堆
            for (int i = k; i < input.length; i++) {
                if (input[i] < heap[1]) {
                    heap[1] = input[i];
                    heapFromTop(1, k, heap);
                }
            }
            System.out.print("heap:");
            for (int i = 1; i <= k; i++) {
                System.out.print(heap[i] + ",");
            }
            System.out.println("");
            //堆排序
            ArrayList<Integer> list = new ArrayList<>();
            for (int num = k; num >= 1; num--) {
                list.add(0, heap[1]);
                heap[1] = heap[num];
                heapFromTop(1, num, heap);
            }
            return list;
        }
    }

    private static void heap(int k, int[] heap) {
        for (int i = k / 2; i >= 1; i--) {
            int left = i * 2;
            int right = i * 2 + 1;
            int tempIndex;
            if (right > k) {
                tempIndex = left;
            } else {
                tempIndex = heap[left] > heap[right] ? left : right;
            }
            if (heap[tempIndex] > heap[i]) {
                int temp = heap[tempIndex];
                heap[tempIndex] = heap[i];
                heap[i] = temp;
                heapFromTop(tempIndex, k, heap);
            }
        }
    }

    private static void heapFromTop(int top, int k, int[] heap) {
        for (int l = top; l <= k / 2; l++) {
            int left = l * 2;
            int right = l * 2 + 1;
            int tempIndex;
            if (right > k) {
                tempIndex = left;
            } else {
                tempIndex = heap[left] > heap[right] ? left : right;
            }
            if (heap[tempIndex] > heap[l]) {
                int temp = heap[tempIndex];
                heap[tempIndex] = heap[l];
                heap[l] = temp;
            } else {
                break;
            }
        }
    }
}

堆的知识

1. 堆的定义：堆是一颗每个节点的左右孩子都小于（小顶堆）或者大于（大顶堆）根节点的完全二叉树；
2. 堆的存储：由于是一颗完全二叉树，所以适合使用数组来存储，一般将堆顶元素存储在数组的下标为1的位置，这样第i个节点的左右孩子的位置就是：i*2，i*2+1；最后一颗叶子节点的位置是：size/2。
3. 初始化堆：从最后一颗非叶子节点开始自底向上堆化，如果有节点交换，则对子节点进行自顶向下堆化

 private static void heap(int k, int[] heap) {
        for (int i = k / 2; i >= 1; i--) {
            int left = i * 2;
            int right = i * 2 + 1;
            int tempIndex;
            if (right > k) {
                tempIndex = left;
            } else {
                tempIndex = heap[left] > heap[right] ? left : right;
            }
            if (heap[tempIndex] > heap[i]) {
                int temp = heap[tempIndex];
                heap[tempIndex] = heap[i];
                heap[i] = temp;
                heapFromTop(tempIndex, k, heap);
            }
        }
    }

 private static void heapFromTop(int top, int k, int[] heap) {
        for (int l = top; l <= k / 2; l++) {
            int left = l * 2;
            int right = l * 2 + 1;
            int tempIndex;
            if (right > k) {
                tempIndex = left;
            } else {
                tempIndex = heap[left] > heap[right] ? left : right;
            }
            if (heap[tempIndex] > heap[l]) {
                int temp = heap[tempIndex];
                heap[tempIndex] = heap[l];
                heap[l] = temp;
            } else {
                break;
            }
        }
    }

4. 更新堆：替换堆顶元素，然后自顶向下堆化

 heap[1] = input[i];
 heapFromTop(1, k, heap);

5. 堆排序：堆顶元素有序，先获取堆顶元素，然后用最后一个节点的元素替换堆顶元素，堆大小-1，自顶向下堆化，再次获取堆顶元素，直到堆中只有一个元素

 ArrayList<Integer> list = new ArrayList<>();
  for (int num = k; num >= 1; num--) {
            list.add(0, heap[1]);
            heap[1] = heap[num];
            heapFromTop(1, num, heap);
   }