SW개발/코딩테스트

[Leetcode]Kth Largest Element in an Array

https://leetcode.com/problems/kth-largest-element-in-an-array/description/

 

LeetCode - The World's Leading Online Programming Learning Platform

Level up your coding skills and quickly land a job. This is the best place to expand your knowledge and get prepared for your next interview.

leetcode.com

 

문제 분석

주어진 배열에서 k번째로 큰 숫자를 구하는 문제입니다.

 

처음 시도한 답안 - 실패

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        
        def quick_sort(nums, start, end):
            if start >= end: # 원소가 1개인 경우 정렬할 필요가 없음
                return

            pivot = start # 피봇은 임의상 첫번째 원소로 지정
            left = start + 1 # 왼쪽값은 피봇 바로 다음 원소로 지정, left는 피봇보다 큰 값을 찾아가는 용도
            right = end # right는 피봇보다 작은 값을 찾아가는 용도

            while left <= right: # left 와 right가 엇갈리면 탈출
                while left <= right and nums[left] <= nums[pivot]: # 피봇보다 큰 데이터를 찾을 때까지 left 증가, left는 right보다 작을 때까지
                    left += 1
                while right > start and nums[right] >= nums[pivot]: # 피봇보다 작은 데이터를 찾을 때까지 rigth 감소, right는 start보다 클 때까지
                    right -= 1

                if left > right: # 엇갈린 상태라면
                    nums[right], nums[pivot] = nums[pivot], nums[right] # 작은 값과(right) 피봇 값을 바꿈, 엇갈린 경우라면 left보다 right가 작은 값임
                else: # 엇갈리지 않았다면
                    nums[left], nums[right] = nums[right], nums[left] # 작은 값(left)과 큰 값(right)을 바꿈

                quick_sort(nums, start, right-1) # 왼쪽
                quick_sort(nums, right+1, end) # 오른쪽

        quick_sort(nums, 0, len(nums)-1)

        return nums[-k]

접근 방법

  1. 처음에는 quicksort를 이용해 정렬후 몇번째로 큰 원소인지 구하고자 했습니다.
  2. 피봇은 첫번째 원소로 지정합니다. (혹은 start~end 사이의 랜덤값)
  3. left는 시작값보다 1 큰 지점을 선택합니다.
  4. right는 끝값으로 선택합니다.
    1. left와 right가 엇갈릴때까지 반복합니다.
    2. 피봇값보다 큰 값을 찾을때까지 left를 오른쪽으로 이동합니다.
    3. 피봇값보다 작은 데이터를 찾을때까지 right를 왼쪽으로 이동합니다.
    4. 만약 left와 right가 엇갈린 상태라면
      1. 작은 값과(right)와 피봇값을 스왑합니다.
    5. 엇갈리지 않은 상태라면
      1. 작은 값(left)와 큰 값(right)을 바꿉니다.
    6. 이후 재귀적으로 왼쪽과 오른쪽에 대해 quick 정렬을 수행합니다.

처음에는 quick 정렬을 수행하려고 했으나 시간초과로 인해서 풀이할 수 없었습니다.

퀵 정렬의 경우 worst 시간복잡도가 O(n^2)이 되기 때문입니다. 따라서 이 부분을 최적화하려고 pivot 값을 랜덤이나, 중간값, 처음값등 다양하게 바꿔가며 시도했지만 시간초과를 벗어날 수 없었습니다. 

그래서 파이썬의 heapq를 통한 솔루션을 참고해보았습니다.

 

두번째로 시도한 답안

import heapq


class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:
        # heapq 모듈을 사용해 k번째까지 큰 숫자들을 구해옴. 가장 마지막 요소가 k번째로 큰수가 됨
        return heapq.nlargest(k, nums)[-1]

접근 방법

  1. 파이썬의 heapq 모듈을 사용합니다.
  2. heapq.nlargest 함수를 사용해 리스트에서 k번째까지 큰 숫자들을 구해옵니다.
  3. 가장 마지막에 담긴 값이 k번째로 큰수입니다.

heapq 모듈을 사용해 편하게 풀 수 있었습니다. 조금 찾아보니 quick_sort 대신 quick select 방식으로도 풀이가 가능한 것 같은데 마찬가지로 타임아웃이 난다는 discussion이 많이 있어서 시도해보지는 않았습니다.

 

728x90