본문 바로가기
IT/알고리즘

[코딩 인터뷰] 배열에서 k번째로 큰 요소 찾기

by 네야나라 2024. 3. 28.

문제 설명:

정수 배열 nums와 정수 k가 주어질 때, 배열에서 k번째로 큰 요소를 반환하는 함수를 작성하세요. nums에는 중복된 요소가 포함될 수 있으며, k번째로 큰 요소를 찾을 때 중복을 고려합니다. 예를 들어, 배열이 [3,2,3,1,2,4,5,5,6]이고 k=4라면, 배열에서 네 번째로 큰 요소는 4입니다.

입력:

  • nums: 정수 배열
  • k: 찾고자 하는 순서의 정수

출력:

  • 배열에서 k번째로 큰 요소

예시:

입력: nums = [3,2,1,5,6,4], k = 2   출력: 5

입력: nums = [3,2,3,1,2,4,5,5,6], k = 4   출력: 4

 

풀이 및 해설 (Python):

 

퀵 정렬의 변형인 퀵 선택(Quick Select) 알고리즘을 사용하여, 전체 배열을 정렬하지 않고도 k번째로 큰 요소를 효율적으로 찾는 방법

def findKthLargest(nums, k):
    # 퀵 선택을 위한 헬퍼 함수
    def quickSelect(l, r):
        pivot, p = nums[r], l
        for i in range(l, r):
            if nums[i] >= pivot:
                nums[i], nums[p] = nums[p], nums[i]
                p += 1
        nums[p], nums[r] = nums[r], nums[p]
        
        # k번째로 큰 요소를 찾았는지 확인
        if p > k - 1:
            return quickSelect(l, p - 1)
        elif p < k - 1:
            return quickSelect(p + 1, r)
        else:
            return nums[p]
    
    return quickSelect(0, len(nums) - 1)

# 예시 실행
print(findKthLargest([3,2,1,5,6,4], 2))  # 출력: 5
print(findKthLargest([3,2,3,1,2,4,5,5,6], 4))  # 출력: 4

 

이 코드는 다음과 같이 작동합니다:

퀵 선택 알고리즘의 핵심 함수인 quickSelect를 정의합니다. 이 함수는 주어진 배열의 부분집합 [l, r]에 대해 피벗을 기준으로 배열을 두 부분으로 분할하고, k번째로 큰 요소가 어느 쪽에 있는지 결정한 후, 해당 부분집합에 대해 재귀적으로 함수를 호출합니다.
피벗을 배열의 마지막 요소로 선택하고, 배열을 피벗보다 크거나 같은 요소와 작은 요소로 분할합니다. 이 때, 피벗 자신도 적절한 위치에 배치됩니다.
분할된 후, 피벗의 위치(p)가 k-1과 같으면 k번째로 큰 요소를 찾은 것이므로 해당 요소를 반환합니다. p가 k-1보다 크면, k번째로 큰 요소는 왼쪽 부분집합에 있으므로 왼쪽 부분집합에 대해 재귀 호출을 합니다. p가 k-1보다 작으면 오른쪽 부분집합에 대해 재귀 호출을 합니다.

반응형