TOP K问题(大顶堆、分治思想)

package com.heu.wsq.niuke.top200;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.PriorityQueue;

/**
 * TOP K问题
 * @author wsq
 * @date 2021/5/8
 * 题目描述
 * 给定一个数组，找出其中最小的K个数。例如数组元素是4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4。如果K>数组的长度，那么返回一个空的数组
 *
 * 示例1
 * 输入：
 *      [4,5,1,6,2,7,3,8],4
 * 输出：
 *      [1,2,3,4]
 */
public class GetLeastNumbers {
    /**
     * 大顶堆
     * 使用大顶堆保存k个值，下一次新来的值v需要跟堆顶比较，
     * 如果堆顶元素大于v，则将新此时的堆顶弹出，将v加入到大顶堆中
     * 如果堆顶元素小于v，则直接抛弃v
     * @param input
     * @param k
     * @return
     */
    public ArrayList<Integer> GetLeastNumbers_Solution(int [] input, int k) {
        if(input == null || input.length < k){
            return new ArrayList<Integer>();
        }
        PriorityQueue<Integer> pq = new PriorityQueue<Integer>(new Comparator<Integer>(){
            public int compare(Integer o1, Integer o2){
                return o2 - o1;
            }
        });
        int size = 0;
        for(int i = 0; i < input.length; i++){
            pq.offer(input[i]);
            size++;
            if(size > k){
                pq.poll();
            }
        }
        ArrayList<Integer> ans = new ArrayList<>();
        while(!pq.isEmpty()){
            ans.add(pq.poll());
        }
        return ans;
    }

    /**
     * 快速排序的分治思想
     * 通过partition将数组分为两部分，
     * 索引p之前：数值小于索引p对应的元素值
     * 索引p之后：数值大于索引p对应的元素值
     * 如果p+1正好为k个元素，则索引p就是要寻找的位置，p之前的正好k个元素
     * @param input
     * @param k
     * @return
     */
    public ArrayList<Integer> GetLeastNumbers_Solution2(int [] input, int k) {
        if (input == null || input.length < k || k==0){
            return new ArrayList<>();
        }
        int left = 0;
        int right = input.length - 1;
        int p = 0;
        while(left <= right){
            p = partition(input, left, right);
            if(p + 1 == k){
                break;
            }else if(p + 1 > k){
                right = p - 1;
            }else if(p + 1 < k){
                left = p + 1;
            }
        }
        ArrayList<Integer> ans = new ArrayList<>();
        for(int i = 0; i <= p; i++){
            ans.add(input[i]);
        }
        return ans;
    }

    private int partition(int[] input, int left, int right) {
        int v = input[left];
        int i = left + 1;
        int j = right;
        while(true){
            while(i <= right && input[i] <= v){
                i++;
            }
            while(j >= left && input[j] > v){
                j--;
            }
            if(i >= j){
                break;
            }
            swap(input, i, j);
        }
        swap(input, left, j);
        return j;
    }

    private void swap(int[] input, int i, int j) {
        int tmp = input[i];
        input[i] = input[j];
        input[j] = tmp;
    }
}