Skip to main content

Kth Largest Element in an Array

문제

정수 배열 nums와 정수 k가 주어지면 배열에서 k번째로 큰 요소를 반환합니다. 이는 정렬된 순서에서 k번째로 큰 요소이지 k번째 고유 요소가 아닙니다. 정렬하지 않고 해결하세요.

Example 1:

Input: nums = [3,2,1,5,6,4], k = 2
Output: 5

Example 2:

Input: nums = [3,2,3,1,2,4,5,5,6], k = 4
Output: 4

나의 풀이

from heapq import nlargest

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

nlargest 함수는 입력된 리스트를 내림차순으로 가장 큰 n개의 요소를 반환합니다. 예를 들어 주어진 입출력 예시처럼 [3, 2, 1, 5, 6, 4] 배열이 nums로 주어지고 2번째로 큰 수를 출력하기 위해서 nlargest(2, [3, 2, 1, 5, 6, 4]) 함수를 실행하여 반환받는 리스트는 [6, 5]가 됩니다. nlargest 함수가 반환하는 리스트 중 가장 마지막 수가 k번째로 큰 수 이므로 리스트의 마지막 요소를 반환합니다.

다른 사람의 풀이1

class Solution(object):
def findKthLargest(self, nums, k):
"""
:type nums: List[int]
:type k: int
:rtype: int
"""
min_value = min(nums) # 배열 요소 중 최솟값
max_value = max(nums) # 배열 요소 중 최댓값

count = [0] * (max_value - min_value + 1)

for num in nums:
count[num - min_value] += 1
# 각 인덱스는 요소의 개수를 나타내며
# [0, 0, 1, 0, 0, 0]
# 1 2 3 4 5 6 -> 요소 3이 존재하므로 개수 증가
# [0, 1, 1, 0, 0, 0]
# 1 2 3 4 5 6 -> 요소 2가 존재하므로 개수 증가
# [1, 1, 1, 0, 0, 0]
# [1, 1, 1, 0, 1, 0]
# [1, 1, 1, 0, 1, 1]
# [1, 1, 1, 1, 1, 1]

# k번째로 큰 수를 구하기 위해서 가장 마지막부터 순회
# count 배열에서 가장 마지막 인덱스가 가장 큰 요소의 개수를 의미
remain = k
for num in range(len(count) - 1, -1, -1):
# 가장 마지막 큰 요소부터 차감한다면 remain이 0이 되었을 때
# k번째로 큰 요소의 인덱스를 가르키게 되고 이를 요소 변환하여 반환
remain -= count[num]
if remain <= 0:
# num - min_value로 인덱스를 표시했다면
# num + min_value로 요소로 변환하여 요소 반환
return num + min_value
return -1

max_value - min_value + 1로 배열의 길이를 설정하는 이유는 계수 정렬에서는 최솟값과 최댓값 사이의 모든 요소를 배열로 나타낼 수 있어야 합니다. 즉 nums에 포함된 모든 정수 값을 인덱스로 사용할 수 있게 하기 위해서 위와 같은 배열의 길이로 초기화합니다. 계수 정렬은 이렇게 모든 요소를 나타낼 수 있도록 초기화하므로 많은 메모리 공간을 필요로 합니다. 만약 최댓값과 최솟값의 간격이 그만큼 크다면 두 간격 만큼 메모리 효율이 떨어질 수 있습니다.

다른 사람의 풀이2

import heapq

class Solution(object):
def findKthLargest(self, nums, k):
"""
:type nums: List[int]
:type k: int
:rtype: int
"""
heap = []
for num in nums:
heapq.heappush(heap, num)
if len(heap) > k:
heapq.heappop(heap)

return heap[0]

최소 힙을 사용하여 힙의 루트에 가장 작은 숫자가 위치하는 원리를 사용한 풀이이며 k번째로 큰 수를 찾기 위해 힙의 크기를 k로 유지하면서 모든 숫자를 힙에 넣을 경우 힙에 남는 요소 중 가장 첫 번째 원소가 k번째로 큰 요소가 남게 됩니다.