215. Kth Largest Element in an Array-Python

利用Python最大堆解决问题。

Find the kth largest element in an unsorted array. Note that it is the kth largest element in the sorted order, not the kth distinct element.

For example,
Given [3,2,1,5,6,4] and k = 2, return 5.

Note: 
You may assume k is always valid, 1 ≤ k ≤ array's length.

Credits:
Special thanks to @mithmatt for adding this problem and creating all test cases.

Code

class Solution(object):

    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        #first build heap
        n=len(nums)
        i=(n-1)/2
        while i>=0:
            self.buildHeap(nums,i)
            i-=1    
        i=1   
        while i<k:
            # first element swaps to last element   
            nums[0]=nums[n-i] #override max element
            nums=nums[0:n-i] 
            self.buildHeap(nums,0) 
            i+=1                    
        return nums[0]

    def buildHeap(self,nums,i):
        init=i
        leftIndex=self.getLeft(nums,i)
        rightIndex=self.getRight(nums,i)
        if leftIndex!=-1 and nums[i]<nums[leftIndex]:                
            i=leftIndex
        if rightIndex!=-1 and nums[i]<nums[rightIndex]:
            i=rightIndex        
        if i!=init:#check if i changes
            tmp=nums[init]#swap
            nums[init]=nums[i]
            nums[i]=tmp
            self.buildHeap(nums,i)

    # get left index for i index
    def getLeft(self,nums,i):
        """return left element for i
        if no exists, return -1
        """
        return -1 if 2*i+1>=len(nums) else 2*i+1
        # get right index for i index
    def getRight(self,nums,i):
        """return right element for i
        if no exists, return -1
        """
        return -1 if 2*i+2>=len(nums) else 2*i+2

版权声明:本文为daigualu原创文章,遵循 CC 4.0 BY-SA 版权协议,转载请附上原文出处链接和本声明。
THE END
< <上一篇
下一篇>>