面试之巧解TopK

正文

问题描述：

从arr[1, n]这n个数中，找出最大的k个数，这就是经典的TopK问题。

栗子：

从arr[1, 12]={5,3,7,1,8,2,9,4,7,2,6,6} 这n=12个数中，找出最大的k=5个

方法一：k次冒泡排序

我们知道一次冒泡排序可以将一个数‘冒’到目的位置
那么我们只要调用K次冒泡排序即可

方法二：堆排序

朴实无华的手写堆排

class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        #建立最大堆 这个堆的大小是k [0,k-1]
        def heapfy(heap,length,i):
            maxIndex=i
            left=2*i+1
            right=2*i+2
            if left<length and heap[maxIndex]<heap[left]:
                maxIndex=left
            if right<length and heap[maxIndex]<heap[right]:
                maxIndex=right
            if maxIndex!=i:
                heap[i],heap[maxIndex]=heap[maxIndex],heap[i]
                heapfy(heap,length,maxIndex)
        if k==0:return []
        #前k个
        heap=arr[:k]
        for i in range(k//2-1,-1,-1):
            heapfy(heap,k,i)
        #后续
        for i in range(k,len(arr)):
            if arr[i]<heap[0]:
                heap[0]=arr[i]
                heapfy(heap,k,0)
        #用建好的堆进行堆排序
        for i in range(k-1,-1,-1):
            heap[0],heap[i]=heap[i],heap[0]
            heapfy(heap,i,0)
        return heap

方法三：减治法+快排partition

在学习快排的时候，调用一次partition能够将一个数(pivot)安置到最终的位置
那么如果我们用一个第k的数作为pivot，就只要调用一次partition就能得到最终答案！！
前半段就是TopK，后半段不需要管

所以我们是思路是用减治法找到第k大数的下标K_th，
然后用partition处理一遍，即可得到答案

调用减治法找Kth的时间是O(n),调用partition的时间是O(n)
空间是O(1)
完美！

下附 Python和Java的代码
Java代码

class Solution {
    public int partition(int[] arr,int left,int right){
        //首先O(n)扫描
        int pivot=arr[left];
        int location=left;
        for(int i=left+1;i<=right;i++){
            //加不加等号都可以
            if (arr[i]<=pivot){
                swap(arr,++location,i);
            }
        }
        swap(arr,location,left);
        return location;
    }
    public int[] getLeastNumbers(int[] arr, int k) {
        if(k==0)
            return new int[0];
        return quickSearch(arr,0,arr.length-1,k-1);

    }
    public int[] quickSearch(int[]arr,int left,int right,int k){
        int index=partition(arr,left,right);
        if (index==k)
            return Arrays.copyOf(arr, k+1);
        else if (index > k)
            return quickSearch(arr,left,index-1,k);
        else
            return quickSearch(arr,index+1,right,k);
    }
    public void swap(int[] arr,int i,int j){
        int temp=arr[i];
        arr[i]=arr[j];
        arr[j]=temp;
        return ;
    }
}

Python代码

#用快排的思路找到第k小的数即可
class Solution:
    def getLeastNumbers(self, arr: List[int], k: int) -> List[int]:
        def partition(arr,left,right):
            pivot=arr[left]
            location=left #location指向小于等于pivot的数组段的尾部
            for i in range(left+1,right+1): 
                if arr[i]<pivot:
                    location+=1
                    arr[location],arr[i]=arr[i],arr[location]
            arr[left],arr[location]=arr[location],arr[left]
            return location #返回的是下标

        #目的是找第k大的下标
        def quickSearch(arr,left,right,k):
            index=partition(arr,left,right)
            if index==k:
                return arr[:k+1]
            if index>k:
                return quickSearch(arr,left,index-1,k)
            else:
                return quickSearch(arr,index+1,right,k)
        

        return [] if k==0 else quickSearch(arr,0,len(arr)-1,k-1)