Python에서 선형 시간에서 k 번째로 작은 요소를 찾는 프로그램

<시간/>

nums라는 숫자 목록이 있고 또 다른 값 k가 있다고 가정하고 목록에서 k번째(0부터 시작) 가장 작은 요소를 찾아야 합니다. 평균 O(n) 시간 안에 이 문제를 풀어야 합니다.

따라서 입력이 nums =[6, 4, 9, 3, 1] k =2와 같으면 출력은 4가 됩니다. 정렬 후 목록은 [1, 3, 4, 6, 9]와 같을 것입니다. , k번째로 작은 요소는 4입니다.

이 문제를 해결하기 위해 다음 단계를 따릅니다. −

maxHeap :=새로운 빈 힙
0~k 범위의 i에 대해
- nums[i]를 maxHeap에 삽입
k + 1 범위에서 숫자 - 1 크기의 i에 대해
- nums[i]> maxHeap[0]이면
  - maxHeap에서 상단 삭제
  - nums[i]를 maxHeap에 삽입
maxHeap[0] 반환

예시

이해를 돕기 위해 다음 구현을 살펴보겠습니다. −

from heapq import heappop, heappush
def solve(nums, k):
   maxHeap = []
   for i in range(k + 1):
      heappush(maxHeap, -nums[i])
   for i in range(k + 1, len(nums)):
      if nums[i] < -maxHeap[0]:
         heappop(maxHeap)
         heappush(maxHeap, -nums[i])
   return -maxHeap[0]

nums = [6, 4, 9, 3, 1]
k = 2
print(solve(nums, k))

입력

[6, 4, 9, 3, 1], 2