面試之巧解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)