Find Kth-Largest: Quick Select Algorithm Exceeds Max Time with Large K

Question

I'm working on a Leet Code Question. Finding the Kth-largest element in an array.

This is the link to the question: https://leetcode.com/problems/kth-largest-element-in-an-array/

This is my solution, I tried to solve it via the Quick Select Algorithm. However, when given a very large array, and a large K (50000). My algorithm exceeds the maximum time.

I'm confused on where in the algorithm is causing this slow down or if i made some logic mistakes, hope to get some advice on this, thanks in advance.

class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """

        n = len(nums)
        random.shuffle(nums)

        return self.quickselect(nums, 0, n-1, n-k)


    def quickselect(self, nums, l, r, k):

        piv = self.partition(nums, l, r)

        if k < piv:
            return self.quickselect(nums, l, piv-1, k)
        elif k > piv:
            return self.quickselect(nums, piv+1, r, k)
        else:
            return nums[piv]


    def partition(self, nums, l, r):
        
        piv = r
        i = l

        for j in range(l, r):

            if nums[j] <= nums[piv]:
                nums[i], nums[j] = nums[j], nums[i]
                i += 1

        nums[i], nums[piv] = nums[piv], nums[i]
        return i

I first thought it's a problem with running into the worst case, since I'm using a fixed end pivot. Hence I shuffled the array. However, this didn't work. The algorithm can handle large arrays with a small K, it is only when K is large as well, it times out.

Answer 1

@greybeard is right. The problem will be hit when nums has only one value, repeated many times, and k is near n/2.

In this case, every time you partition the array, you will only remove one element, leading to quadratic running time.

There are a number of ways of fixing this. One is to do the swap when nums[j] < nums[piv], and do it half of the times when nums[j] == nums[piv].