第29题——最小的k个数

题目：

输入n个整数，找出其中最小的K个数。例如输入4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4,。

思路：

维护一个大小为k的大根堆，遍历这个序列，前k个数进堆，并形成大根堆。后面的数和堆顶比较，若小于对顶，则进堆，并删除堆顶。最终堆中的元素就是最小的k个数了

Java

package nowcoder;

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

public class S29_GetLeastKNumbers {
    public ArrayList<Integer> getLeastKNumbers(int[] input, int k){
        ArrayList<Integer> result = new ArrayList<Integer>();
        int length = input.length;
        if (length == 0 || k <= 0)
            return result;
        //大根堆
        PriorityQueue<Integer> maxHeap = new PriorityQueue<Integer>(k, new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o2.compareTo(o1);
            }
        });
        for (int i=0;i<length;i++){
            if (maxHeap.size() != k){
                maxHeap.offer(input[i]);
            }
            else if (maxHeap.peek() > input[i]){
                maxHeap.remove();
                maxHeap.offer(input[i]);
            }
        }
        result.addAll(maxHeap);
        return result;
    }
    public static void main(String[] args){
        S29_GetLeastKNumbers s29 = new S29_GetLeastKNumbers();
        int[] input = {4,5,1,6,2,7,3,8};
        System.out.println(s29.getLeastKNumbers(input, 4));
    }
}

不利用PriorityQueue，重新写堆排序

package nowcoder;

import java.util.ArrayList;

public class S29_GetLeastKNumbers_Original {
    //调整大根堆
    public void sift(int[] array,int low,int high){
        int i = low;
        int temp = array[i];
        int j = 2*i+1; //i的左孩子
        while (j <= high){
            if (j < high && array[j+1] > array[i])
                j++;
            if (temp < array[j]){
                array[i] = array[j];
                i = j;
                j = 2*i+1;
            }else
                break;
        }
        array[i] = temp;
    }
    public ArrayList<Integer> getLeastKNumbers(int[] input, int k){
        ArrayList<Integer> result = new ArrayList<Integer>();
        if (input.length == 0 || k<= 0)
            return result;
        int[] numbers = new int[k];
        for (int i=0;i<k;i++)
            numbers[i]=input[i];
        //构建容量为k的大根堆
        for (int i=k/2-1;i>=0;i--)
            sift(numbers,i,k-1);
        //后面的数和堆顶比较，比堆顶小就替换堆顶，并重新调整成大根堆
        for (int i=k;i<input.length;i++){
            if (numbers[0] > input[i]){
                numbers[0] = input[i];
                sift(numbers,0,k-1);
            }
        }
        for (int n:numbers)
            result.add(n);
        return result;
    }
    public static void main(String[] args){
        S29_GetLeastKNumbers_Original s29 = new S29_GetLeastKNumbers_Original();
        int[] input = {4,5,1,6,2,7,3,8};
        System.out.println(s29.getLeastKNumbers(input, 4));
    }
}

Python

class GetLeastKNumbers:
    def getLeastKNumbers(self, input, k):
        if len(input) == 0 or k <=0 or k > len(input):
            return []
        #构建容量为k的大根堆
        for i in range(k//2, -1, -1):
            self.sift(input, i, k-1)
        for i in range(k, len(input)):
            if input[i] < input[0]:
                input[0] = input[i]
                self.sift(input, 0, k-1)
        numbers = []
        for i in range(k):
            numbers.append(input[i])
        return numbers
    def sift(self,array,low,high):
        i = low
        temp = array[i]
        j = 2*i+1
        while j <= high:
            if j < high and array[j] < array[j+1]:
                j += 1
            if temp < array[j]:
                array[i] = array[j]
                i = j
                j = i*2+1
            else:
                break
        array[i] = temp

if __name__ == '__main__':
    test = GetLeastKNumbers()
    input = [4,5,1,6,2,7,3,8]
    print(test.getLeastKNumbers(input, 3))

结果：

[3, 2, 1]