215. Kth Largest Element in an Array

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

Kth Largest Element in an Array - LeetCode

Solution

import random

class Solution:
    def findKthLargest(self, nums, k):
        k = len(nums) - k
        return self.quick(nums, k)

    def quick(self, nums, k):
        L = len(nums)
        idx = random.randrange(0, L)
        pivot = nums[idx]
        i, j = 0, L-1

        while(i<j):
            while(i < L and nums[i] <= pivot):
                i += 1
            while(j >= 0 and nums[j] >= pivot):
                j -= 1
            if i<j:
                nums[i], nums[j] = nums[j], nums[i]

        if idx > i:
            nums[idx], nums[i] = nums[i], nums[idx]
            idx = i
        elif idx < j:
            nums[idx], nums[j] = nums[j], nums[idx]
            idx = j

        if k < idx :
            return self.quick(nums[:idx], k)
        elif k > idx :
            return self.quick(nums[idx+1:], k-idx-1)
        else:
            return nums[k]

• sort를 쓰면 편하지만 불필요한 정렬을 하지 않고 풀고 싶었다.
• quick sort의 방식으로 pivot을 기준으로 partition을 나누면 k가 포함되는 한쪽만 정렬해서 최소한으로 정렬하도록 하였다.
• quick의 인자로 partition과 찾아야할 k를 넣는다.
• 랜덤으로 pivot을 고른다. 그렇지 않으면 [3, 3, 3, 4, 3, 3] 과 같은 배열은 정렬할 수 없다.
• 왼쪽에서는 pivot보다 큰 수를 찾고, 오른쪽에서는 pivot보다 작은 수를 찾는다.
• 그렇게 찾았을 때 큰 수가 작은 수보다 왼쪽에 있다면 swap한다.
• 그렇게 반으로 나눴다면 pivot을 중간 자리로 옮긴다.
• pivot의 위치와 k 위치를 비교해서 어느 쪽 partition을 다음에 정렬해야 할지 선택한다.
• 만약 pivot의 위치가 k라면 반환한다.

Another 1

class Solution:
    def findKthLargest(self, nums, k):
        return sorted(nums)[-k]

• 가장 짧은 코드이다.
• 생각보다 그렇게 느리진 않았다.
• 실제 코딩테스트에서는 이렇게 푸는게 시간상 옳다.

Another 2

from random import randint
class Solution(object):
    def findKthLargest(self, nums, k):
        def partition(l, r):
            ri = randint(l, r)
            nums[r], nums[ri] = nums[ri], nums[r]
            for i, v in enumerate(nums[l: r+1], l):
                if v >= nums[r]:
                    nums[l], nums[i] = nums[i], nums[l]
                    l += 1
            return l - 1
        
        l, r, k = 0, len(nums) - 1, k - 1
        while True:
            pos = partition(l, r)
            if pos < k:
                l = pos + 1
            elif pos > k:
                r = pos - 1
            else:
                return nums[pos]

• 함수 안에 함수를 정의해서 nums를 인자로 넣지 않아도 사용할 수 있도록 한다.
• 내림차순으로 정렬하여 k를 딱히 바꾸지 않았다.
• partition 함수는 퀵소트처럼 랜덤한 pivot을 고른다.
• 그리고 그 pivot을 가장 오른쪽으로 보낸다.
• for문으로 왼쪽부터 하나씩 보면서 진행하고, pivot보다 크거나 같다면 현재 자리 i와 nums[l]를 바꾼다.
  • enumerate의 두 번째 인자는 인덱스 생성 시작값을 의미한다.
  • 여기서는 l부터 시작해서 인덱스가 하나씩 증가한다.
• l 을 증가시켜서 partition의 왼쪽에 pivot보다 큰 애들이 하나씩 쌓이도록 한다.
• pivot이 가장 마지막에 있으므로 마지막 변경은 pivot과 이루어진다. 이를 위해서 pivot을 가장 오른쪽으로 보낸 것이다. 아주 센스있다.
• 굉장히 깔끔하고 좋은 코드이다.