剑指 Offer 40. 最小的k个数

package leetcode;

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

public class offer_40 {
    public int[] getLeastNumbers(int[] arr, int k) {
        int[] nums=new int[k];
        if(arr.length<=k) {return arr;}
        if(k<=0) {return nums;}
        
        //建立一个大根堆
        PriorityQueue<Integer> priority=new PriorityQueue<Integer>(new Comparator<Integer>() {
            public int compare(Integer a,Integer b) {
                return b-a;
            }
        });
        for(int i=0;i<k;i++) {
            priority.offer(arr[i]);
        }
        //如果当前节点小于根节点则替换根节点
        for(int i=k;i<arr.length;i++) {
            if(priority.peek()>arr[i]) {
                priority.poll();
                priority.offer(arr[i]);
            }
        }
        //将大根堆中的节点值保存在nums中
        for(int i=0;i<k;i++) {
            nums[i]=priority.poll();
        }
        
        return nums;
    }
    public static void main(String[] args) {
        // TODO Auto-generated method stub
        offer_40 off=new offer_40();
        int[] arr= {3,2,1};
        off.getLeastNumbers(arr, 2);
    }

}