第K大元素

1. Link

2. 描述

有一个整数数组，请你根据快速排序的思路，找出数组中第K大的数。

给定一个整数数组a,同时给定它的大小n和要找的K(1<=K<=n)，请返回第K大的数(包括重复的元素，不用去重)，保证答案存在。

示例1

输入：

[1,3,5,2,2],5,3

复制返回值：

复制

示例2

输入：

[10,10,9,9,8,7,5,6,4,3,4,2],12,3

复制返回值：

复制说明：

去重后的第3大是8，但本题要求包含重复的元素，不用去重，所以输出9

3. 思路

method 1: quickselect
1. Time: O(nlogn) if array is not almost sorted, Otherwise, O(n^2)
2. Space: O(1), inplace
method 2: heapsort
1. Time: O(Nlogk)
2. Space: O(k) for heap storage

4. Coding

QuickSelect （quicksort ） method


class Solution:
    def findKth(self, a, n, K):
        #
        # method 1: quickselect
        # method 2: heapsort
        # input: a, n, k
        #
        #
        if K >n or not a:
            return None
        return self.quickselect(a, 0, len(a)-1, K)
    def quickselect(self, arr, start, end, k):
        #
        #1. partition arr into two part
        # 2. get the pos of pivot
        # 3. goto left or right  array to find the k^th element
        #
        
        #binary search
        while start <=end-1:
            pivot = self.partition(arr, start,end)
            if pivot <k-1:
                start = pivot+1
            elif pivot>k-1:
                end = pivot -1
            else:
                return arr[pivot]
        return arr[end]
        
        
        
    def partition(self, arr, start, end):
        if start >= end:
            return end
        pivot = end
        stored_idx = start # store the first element that is greater than pivot
        pt = start
        while pt <end:
            # since it is Kth largest, so sort it in descending order
            if arr[pt] >= arr[pivot]:
                arr[pt], arr[stored_idx] = arr[stored_idx], arr[pt] 
                stored_idx += 1
            pt += 1
        arr[stored_idx], arr[pivot] =arr[pivot], arr[stored_idx]
        return stored_idx

2. Heap Sort method

class Solution:
    def findKth(self, a, n, K):
        # method2: heap sort
        # 1. build a min-heap using first K element in array
        # 2. use min-heap to pop the minimum element repeatedly until
        #   the heap contain only K leement. Then the last element as top k largest
        #
        #
        # Time: O(klogk + (n-k)logk + logk) = O(nlogk)
        # Build heap: O(klogk) for build heap, while loop O((n-k)logk), O(logk) for return result
        # 
        import heapq
        if K >n and not a:
            return None
        heap = [a[i] for i in range(K)]
        heapq.heapify(heap)
        pt = K
        while pt < n:
            
            heapq.heappush(heap, a[pt])
            heapq.heappop(heap)
            pt += 1
        
        # return the minimum value in the K largest values of array
        # that is TopK largest 
        return heapq.heappop(heap)

PreviousMergeSort Linkedlist NextLinkedList

Last updated 3 years ago

Was this helpful?

class Solution: def findKth(self, a, n, K): # # method 1: quickselect # method 2: heapsort # input: a, n, k # # if K >n or not a: return None return self.quickselect(a, 0, len(a)-1, K) def quickselect(self, arr, start, end, k): # #1. partition arr into two part # 2. get the pos of pivot # 3. goto left or right array to find the k^th element # #binary search while start <=end-1: pivot = self.partition(arr, start,end) if pivot <k-1: start = pivot+1 elif pivot>k-1: end = pivot -1 else: return arr[pivot] return arr[end] def partition(self, arr, start, end): if start >= end: return end pivot = end stored_idx = start # store the first element that is greater than pivot pt = start while pt <end: # since it is Kth largest, so sort it in descending order if arr[pt] >= arr[pivot]: arr[pt], arr[stored_idx] = arr[stored_idx], arr[pt] stored_idx += 1 pt += 1 arr[stored_idx], arr[pivot] =arr[pivot], arr[stored_idx] return stored_idx

class Solution: def findKth(self, a, n, K): # method2: heap sort # 1. build a min-heap using first K element in array # 2. use min-heap to pop the minimum element repeatedly until # the heap contain only K leement. Then the last element as top k largest # # # Time: O(klogk + (n-k)logk + logk) = O(nlogk) # Build heap: O(klogk) for build heap, while loop O((n-k)logk), O(logk) for return result # import heapq if K >n and not a: return None heap = [a[i] for i in range(K)] heapq.heapify(heap) pt = K while pt < n: heapq.heappush(heap, a[pt]) heapq.heappop(heap) pt += 1 # return the minimum value in the K largest values of array # that is TopK largest return heapq.heappop(heap)