程序超出时间限制 - 找到第 K 个最小元素

Question

I am looking at GeeksForGeeks problem Kth smallest element :

Given an array arr[] and an integer K where K is smaller than size of array, the task is to find the Kth smallest element in the given array. It is given that all array elements are distinct.

Expected Time Complexity: O(n)
Expected Auxiliary Space: O(log(n))
Constraints:
1 <= N <= 105
1 <= arr[i] <= 105
1 <= K <= N

My Code:

class Solution:
    def kthSmallest(self,arr, l, r, k):
        '''
        arr : given array
        l : starting index of the array i.e 0
        r : ending index of the array i.e size-1
        k : find kth smallest element and return using this function
        '''
        arr2=arr[:k]
        arr2.insert(0,None)
        for i in range(k//2,0,-1):
            arr2=self.heapify(arr2,i,k-1)
            
        for i in arr[k:]:
            if i <arr2[1]:
                arr2[1]=i
                arr2=self.heapify(arr2,1,k-1)
        return arr2[1]


    def heapify(self,arr, i, r):
        if 2 * i <= r + 1 and arr[2 * i] > arr[i]:
            arr[2 * i], arr[i] = arr[i], arr[i * 2]
            arr = self.heapify(arr, 2 * i, r)
        if 2 * i + 1 <= r + 1 and arr[2 * i + 1] > arr[i]:
            arr[2 * i + 1], arr[i] = arr[i], arr[i * 2 + 1]
            arr = self.heapify(arr, 2 * i + 1, r)
        return arr

I made a sub array of first K elements in the array, and max heapified it. Then for the rest of the elements in the array, if the element is smaller than the first element of the heap, I replaced the top element and then max heapified the top element. I am getting time limit exceeded error. Any idea?

Answer 1

The problem is that your heapify function is not efficient. In the worst case it makes two recursive calls at the same recursion depth. This may even happen at several recursion depths, so that the number of times heapify is called recursively could become quite large. The goal is to have this only call heapify once (at the most) per recursion level.

It should first find the greatest child, and only then determine whether heapify should be called again, and make that single call if needed.

Some other remarks:

• Instead of making heapify recursive, use an iterative solution. This will also save some execution time.

• It is strange to pass k-1 to heapify as last argument, when the last element sits at index k , and so you get the weird comparison <= r + 1 in that function. It is more intuitive to pass k as argument, and work with <= r inside the function.

• As arr is mutated by heapify it is not needed to return it. This is just overhead that is useless.

• 2 * i is calculated several times. It is better to calcuate this only once.

• arr[k:] makes a copy of that part of the list. This is not really needed. You could just iterate over the range and take the corresponding value from the array in the loop.

• It is not clear why the main function needs to get l and r as arguments, since in comments it is explained that l will be 0 and r the index of the last element. But in my opinion, since you get them, you should use them. So you should not assume l is 0,... etc.

• I would use a more descriptive name for arr2 . Why not name it heap ?

Here is the improvement of your code:

class Solution:
    def kthSmallest(self,arr, l, r, k):
        '''
        arr : given array
        l : starting index of the array i.e 0
        r : ending index of the array i.e size-1
        k : find kth smallest element and return using this function
        '''
        heap = [arr[i] for i in range(l, r + 1)]
        heap.insert(0, None)
        for i in range(k//2, 0, -1):
            self.heapify(heap, i, k)
        for i in range(l + k, r + 1):
            val = arr[i]
            if val < heap[1]:
                heap[1] = val
                self.heapify(heap, 1, k)
        return heap[1]

    def heapify(self, arr, i, r):
        child = 2 * i
        while child <= r:
            if child + 1 <= r and arr[child + 1] > arr[child]:
                child += 1
            if arr[child] <= arr[i]:
                break
            arr[child], arr[i] = arr[i], arr[child]
            i = child
            child = 2 * i

Finally, there is a heapq module you can use, which simplifies your code:

from heapq import heapify, heapreplace

class Solution:
    def kthSmallest(self, arr, l, r, k):
        '''
        arr : given array
        l : starting index of the array i.e 0
        r : ending index of the array i.e size-1
        k : find kth smallest element and return using this function
        '''
        heap = [-arr[i] for i in range(l, l + k)]
        heapify(heap)
            
        for i in range(l + k, r + 1):
            val = -arr[i]
            if val > heap[0]:
                heapreplace(heap, val)
        return -heap[0]

The unary minus that occurs here and there is to make the native minheap work as a maxheap.