[LeetCode] 786. K-th Smallest Prime Fraction

문제 출처: https://leetcode.com/problems/k-th-smallest-prime-fraction/

K-th Smallest Prime Fraction - LeetCode

 

문제

1과 소수를 포함하고 있는 정수 배열 arr이 주어진다. arr의 모든 정수는 유일하다. 또한 정수 k가 주어진다.

 

모든 0 <= i < j < arr.length인 i와 j에 대해서 분수 arr[i] / arr[j] 를 고려한다.

 

kth로 가장 작은 분수를 반환하라. 답은 크기가 2인 배열이여야 하고, answer[0] == arr[i] 그리고 answer[1] == arr[j]이어야 한다.

 

예제

Input: arr = [1,2,3,5], k = 3
Output: [2,5]

Explanation: The fractions to be considered in sorted order are:
1/5, 1/3, 2/5, 1/2, 3/5, and 2/3.
The third fraction is 2/5.

Input: arr = [1,7], k = 1
Output: [1,7]

 

풀이

문제에 가장 쉽게 접근할 수 있는 방법으로 Brute Force가 있다. 단순하게 표현할 수 있는 모든 분수 표현을 만들어 배열로 저장한 다음, 정렬 하거나 우선순위 큐를 이용해서 k 번째로 작은 값을 찾는 것이다.

 

우선순위 큐를 이용할 경우 아래와 같은 코드를 작성할 수 있다.

(JavaScript의 경우 heap을 별도로 구현해야 하기 때문에 문제를 해결하는 방식에 조금 더 집중하기 위해서 Pyhon을 이용함)

# Python Code
import heapq

class Solution:
    def kthSmallestPrimeFraction(self, arr: List[int], k: int) -> List[int]:
        frac = [
            (arr[i] / j, arr[i], j) 
            for i in range(len(arr))
            for j in arr[i+1:]
        ]
        heapq.heapify(frac)
        for _ in range(k):
            _, ans_i, ans_j = heapq.heappop(frac)
        return [ans_i, ans_j]

이 접근 방식의 시간 복잡도를 생각해보면, 가능한 분수 표현을 만드는데 O(N²) 이기 때문에 전체 시간 복잡도는 O(N²) 이다. 실제 실행 시간도 2104ms로 개선의 필요가 있다.

 

다음으로 생각해 볼 수 있는 접근 방법은 이진탐색(Binary Search)이다.

 

0 <= i < j < arr.length 의 조건에 의해서 arr[i] / arr[j]로 만들어지는 분수는 0 ~ 1 사이의 값을 갖는다. 이 범위에 대해서 이진 탐색을 수행할 것이다.

 

이진 탐색 알고리즘은 다음과 같은 순서로 동작한다(정렬되어 있는 배열에 대해서).

1. lo = 0, hi = arr.length-1
2. lo, hi의 중간 값(mid)을 계산한다.
3. 탐색하려고 하는 배열에서 mid를 기준으로 왼쪽과 오른쪽 서브배열로 나눈다.
4. 찾고자 하는(target)과 mid를 비교하여 다시 탐색해야 하는 범위가 왼쪽인지 오른쪽인지 판단하여, 2~3 과정을 반복한다.

위 내용을 현재 문제에 맞도록 수정해보자.

1. lo = 0, hi = 1
2. lo, hi의 중간 값(mid)을 계산한다.
3. 분수 표현 중, mid 값보다 작은 분수의 수를 카운트한다.
   • mid 값 보다 작은 분수 표현 중에서 가장 큰 표현을 따로 저장
4. 3에서 카운트한 값(count)과 k를 비교하여,
   • k == count: 저장되어 있는 분수 표현을 만드는 값들을 반환
   • k < count: hi = mid로 설정하여 재탐색 수행
   • k > count: lo = mid로 설정하여 재탐색 수행
5. 4에서 k == count를 찾을 때까지 2~4를 반복

위 내용을 코드로 옮기기 전에 시간 복잡도를 먼저 생각해보자.

 

기본적으로 이진탐색은 k개의 요소에 대해서 알고리즘을 수행할 때, O(log k) 의 시간 복잡도를 갖는다. 여기에 3번 내용을 고려하게 되면, O(N² * log k) 라는 시간 복잡도를 갖게 될 것이라 생각할 수 있다. mid 보다 작은 값의 개수를 찾기 위해 최악의 경우 모든 값을 분수로 만들어봐야 하기 때문이다. O(N²) 과 비교하게 되면 아래의 그림과 같다.

시간 복잡도 비교(검정: n^2, 빨강: N^2 * logN, 보라: N) 

위 그림은 N2 (검정), N² * logN (빨강), N (보라)을 그래프로 그린 것이다. O(N² * logN) 이 O(N2) 보다 좋아 보이기는 하지만, 그 정도가 월등해 보이지는 않는다.  

 

그렇다면 실제 실행시간은 어떻게 나올까. 이를 확인하기 위해서 직접 코드를 작성해보면 아래와 같다.

class Solution:
    def kthSmallestPrimeFraction(self, arr: List[int], k: int) -> List[int]:
        N = len(arr)
        
        lo = 0.0
        hi = 1.0
        
        while lo < hi:
            mid = lo + (hi - lo) / 2
            
            tmp_max_frac = 0
            count = 0
            ans = [0, 0]
            for i in range(N-1):
            	# 분자를 고정시킨 상태에서 분모 값을 변경하며 mid보다 작은 값이 나오는 위치 탐색
                j = i + 1
                while j < N and arr[i] > mid * arr[j]: 
                    j += 1
                
                # 위에서 찾아진 위치를 토대로 몇개의 분수 표현이 존재하는지 카운트
                count += (N - j)
                
                if j == N:
                    break
                
                tmp_frac = arr[i] / arr[j]
                if tmp_max_frac < tmp_frac:
                    tmp_max_frac = tmp_frac
                    ans = [arr[i], arr[j]]
                    
            if count == k:
                return ans
            elif count < k:
                lo = mid
            else:
                hi = mid
        
        return []

위 코드를 제출했을 때, 772ms의 실행시간을 보여주었다. 기존의 방식보다 60% 이상 빨라진 속도이다. 이는 아마도 mid보다 작은 분수 표현을 계산하기 위해서 N2의 연산을 수행하는 일이 자주 일어나지 않기 때문에, 시간 복잡도에서 보이는 것보다 더 드라마틱한 성능 개선이 이루어진 것이라 추측한다.

 

※ 아래 깃 허브의 경우, 이진 탐색을 사용한 방식의 JavaScript 코드이다.

Github: 786.js

GitHub - opwe37/Algorithm-Study