1. 程式人生 > >top k問題python解

top k問題python解

題目來源:Leetcode 215 陣列中第K個最大的元素

題目描述:

在未排序的陣列中找到第 k 個最大的元素。請注意,你需要找的是陣列排序後的第 k 個最大的元素,而不是第 k 個不同的元素。

示例 1:

輸入: [3,2,1,5,6,4] 和 k = 2

輸出: 5

1.排序方法

複雜度為O(nlogn),可以直接呼叫python的排序函式做(用時52ms):

class Solution:
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        nums.sort()
        return nums[-k]
        

自己手寫快排(用時284ms,咋這麼慢。。。):

class Solution:
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        import random
        self.quickSort(nums, 0, len(nums) - 1)
        return nums[-k]
    
    
    def partation(self, nums, l, r):
        """隨機partation"""
        change_index = random.randint(l, r)
        nums[change_index], nums[r] = nums[r], nums[change_index]
        num = nums[r]
        less_index = l - 1
        more_index = r + 1
        p = l
        while p < more_index:
            if nums[p] < num:
                less_index += 1
                nums[p], nums[less_index] = nums[less_index], nums[p]
                p += 1
            elif nums[p] == num:
                p += 1
            else:
                more_index -= 1
                nums[p], nums[more_index] = nums[more_index], nums[p]
                
        equal_area = [less_index + 1, more_index - 1]
        return equal_area
    
    def quickSort(self, nums, l, r):
        if r - l < 1:
            return
        else:
            eq_area = self.partation(nums, l, r)
            self.quickSort(nums, l, eq_area[0] - 1)
            self.quickSort(nums, eq_area[1] + 1, r)

2.選擇方法

定義一個存放最大值的集合set,每遍歷一次陣列,找到一個不在集合中的最大數,返回第k次遍歷結果,這個方法超過時了。

def findKthLargest(nums, k):
    """
    :type nums: List[int]
    :type k: int
    :rtype: int
    """
    num_set = set()
    p = 0
    while p < k:
        cur_max = -float('inf')
        for i in nums:
            if i not in num_set and i > cur_max:
                cur_max = i
                count = 0
            if i == cur_max:
                count += 1
        num_set.add(cur_max)
        p += count
    return cur_max

3.堆

建立一個包括k個數小根堆,遍歷陣列,如果大於小根堆的堆頂,替換,重新構建小根堆。

python自帶的堆實現(52ms):

def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        import heapq
        return heapq.nlargest(k, nums)[-1]

不呼叫庫,自己實現(112ms):

主要解決兩個問題:建立堆,堆頂元素改變維

def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        #先將前k個數變成小根堆
        nums = self.heapInsert(nums, k)
        for i in range(k, len(nums)):
            if nums[i] > nums[0]:
                nums = self.heapify(nums, nums[i], k)
        return nums[0]     
        

def heapInsert(self, nums, k):
        """
        將列表轉變為小根堆
        """
        for index in range(1, k):
            while nums[(index - 1) // 2] > nums[index] and index > 0:
                nums[(index - 1) // 2], nums[index] = \
                nums[index], nums[(index - 1) // 2]
                index = (index - 1) // 2
        return nums

def heapify(self, nums, new_val, k):
        """
        小根堆nums的堆頂變成new_val,重新生成小根堆
        """
        head = 0
        nums[head] = new_val
        l = 1
        while l < k:
            r = l + 1
            if r >= k:
                small = l
            else:
                if nums[l] <= nums[r]:
                    small = l
                else:
                    small = r
        
            if nums[head] < nums[small]:
                break
            nums[head], nums[small] = nums[small], nums[head]
            head = small
            l = head * 2 + 1
        return nums

4.快速選擇

利用快排的思想(這個用了3000多ms,不知道為啥這麼慢。。)

class Solution:
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        #這裡得到的是最大的k個數字,可能是不是有序的
        topk_nums = sorted(self.partation(nums, 0, len(nums) - 1, k))
        return topk_nums[-k]
    
    def partation(self, nums, left, right, k):
        """
        大->中->小
        """
        if right <= left:
            return nums[:k]
        l = left - 1
        r = right + 1
        num = nums[right]
        p = left
        while p < r:
            if nums[p] == num:
                p += 1
            elif nums[p] > num:
                l += 1
                nums[p], nums[l] = nums[l], nums[p]
                p += 1
            else:
                r -= 1
                nums[r], nums[p] = nums[p], nums[r]
        if k <= l:
            return self.partation(nums, left, l, k)
        elif k >= r:
            return self.partation(nums, r, right, k)
        else:
            return nums[:k]

5.插入排序

前k個數字變成一個有序的序列(範圍:0 ~ k - 1,從大到小),遍歷後面的數字,每次遍歷,將其插入到前k個數字中對應的位置(60ms)

class Solution:
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        topk_nums = nums[:k]
        topk_nums.sort()
        for i in range(k, len(nums)):
            if nums[i] > topk_nums[0]:
                insert_index = self.searchInsert(topk_nums, nums[i])
                topk_nums.insert(insert_index, nums[i])
                topk_nums.pop(0)
        return topk_nums[-k]
        
    
    def searchInsert(self, nums, target):
        """
        :type nums: List[int]
        :type target: int
        :rtype: int
        """
        if target <= nums[0]:
            return 0
        if target > nums[-1]:
            return len(nums)

        l, r = 0, len(nums) - 1
        while r >= l:
            mid = (l + r) // 2
            if nums[mid] == target:
                return mid
            elif nums[mid] > target:
                if nums[mid - 1] < target:
                    return mid
                else:
                    r = mid - 1
            else:
                if nums[mid + 1] > target:
                    return mid + 1
                else:
                    l = mid + 1