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?
CodePudding user response:
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
toheapify
as last argument, when the last element sits at indexk
, and so you get the weird comparison<= r 1
in that function. It is more intuitive to passk
as argument, and work with<= r
inside the function.As
arr
is mutated byheapify
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
andr
as arguments, since in comments it is explained thatl
will be 0 andr
the index of the last element. But in my opinion, since you get them, you should use them. So you should not assumel
is 0,... etc.I would use a more descriptive name for
arr2
. Why not name itheap
?
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.