堆的应用:牛客网-求最小的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 GetLeastNumbers_Solution(int[] input, int k) {
        if (k == 0 || input == null || k > input.length) {
            return new ArrayList();
        } 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 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;
            }
        }
    }
  1. 更新堆:替换堆顶元素,然后自顶向下堆化
 heap[1] = input[i];
 heapFromTop(1, k, heap);
  1. 堆排序:堆顶元素有序,先获取堆顶元素,然后用最后一个节点的元素替换堆顶元素,堆大小-1,自顶向下堆化,再次获取堆顶元素,直到堆中只有一个元素
 ArrayList list = new ArrayList<>();
  for (int num = k; num >= 1; num--) {
            list.add(0, heap[1]);
            heap[1] = heap[num];
            heapFromTop(1, num, heap);
   }

你可能感兴趣的:(堆的应用:牛客网-求最小的K个数)