每日一道Leetcode-面试题 17.14. 最小K个数【快速排序】

class Solution:
    def partition(self,nums,l,r):
        # pivot存在最右边
        """
        3 1 5 7 4
        i = -1
        j=0: 3 < 4: i = -1+1 = 0 nums[0] = 3 nums[0] = 3
        j=1: 1 < 4: i = 0+1 = 1 nums[1] = 1 nums[1] = 1
        j=2: 5 > 4
        j=3: 7 > 4
        i + 1 = 2: nums[2] = 4 nums[4] = 4
        """
        pivot = nums[r]
        i = l-1
        for j in range(l,r):
            if nums[j]<=pivot:
                i += 1
                nums[i],nums[j] = nums[j],nums[i]
        nums[i+1],nums[r] = nums[r],nums[i+1]
        return i+1

    def randomized_partition(self,nums,l,r):
        i = random.randint(l,r)
        nums[r],nums[i] = nums[i], nums[r]
        return self.partition(nums,l,r)
    
    def randomized_selected(self,arr,l,r,k):
        # 判断pivot左右两边元素个数与k的关系
        pos = self.randomized_partition(arr,l,r)
        num = pos-l+1
        if k < num:
            self.randomized_selected(arr,l,pos-1,k)
        elif k > num:
            self.randomized_selected(arr,pos+1,r,k-num)

    def smallestK(self, arr: List[int], k: int) -> List[int]:
        if k == 0:
            return []
        self.randomized_selected(arr,0,len(arr)-1,k)
        return arr[:k]