정수 배열 a
가 주어진다. 이 배열에서 k번째로 큰 요소를 찾는 알고리즘을 작성한다.
예를 들어, 1번째로 큰 요소는 배열에서 가장 큰 값이다. 배열에 n개의 요소가 있다면, n번째로 큰 요소는 가장 작은 값이다. 중앙값은 n/2번째로 큰 요소이다.
다음 해결책은 단순한 접근 방식을 취한다. 이 알고리즘은 배열을 먼저 정렬하기 때문에 시간 복잡도가 **O(n log n)**이다. 또한 추가적으로 O(n) 공간을 사용한다.
func kthLargest(a: [Int], k: Int) -> Int? {
let len = a.count
if k > 0 && k <= len {
let sorted = a.sorted()
return sorted[len - k]
} else {
return nil
}
}
kthLargest()
함수는 두 가지 매개변수를 받는다: 정수로 구성된 배열 a
와 k
이다. 이 함수는 k-번째로 큰 요소를 반환한다.
예제를 통해 알고리즘이 어떻게 동작하는지 살펴보자. k = 4
이고 배열이 다음과 같다고 가정한다:
[ 7, 92, 23, 9, -1, 0, 11, 6 ]
처음에는 k번째 큰 요소를 직접 찾을 수 없지만, 배열을 정렬하면 간단해진다. 정렬된 배열은 다음과 같다:
[ -1, 0, 6, 7, 9, 11, 23, 92 ]
이제 a.count - k
인덱스의 값을 찾기만 하면 된다:
a[a.count - k] = a[8 - 4] = a[4] = 9
물론 k번째 작은 요소를 찾고 싶다면 a[k-1]
을 사용한다.
이진 탐색과 퀵 정렬의 아이디어를 결합해 O(n) 시간 복잡도의 해결책을 찾는 기발한 알고리즘이 있다.
이진 탐색은 배열을 반복적으로 반으로 나누어 원하는 값을 빠르게 찾는다. 여기서도 같은 방식을 사용한다.
퀵 정렬 또한 배열을 분할한다. 피벗을 기준으로 더 작은 값은 왼쪽으로, 더 큰 값은 오른쪽으로 이동시킨다. 특정 피벗을 기준으로 분할한 후, 그 피벗 값은 이미 최종 정렬된 위치에 있게 된다. 이를 활용할 수 있다.
작동 방식은 다음과 같다. 무작위로 피벗을 선택하고, 그 피벗을 기준으로 배열을 분할한 후, 이진 탐색처럼 왼쪽 또는 오른쪽 파티션에서만 계속 진행한다. 이 과정을 반복해 k-번째 위치에 오는 피벗을 찾을 때까지 진행한다.
원래 예제를 다시 살펴보자. 다음 배열에서 4번째로 큰 요소를 찾고 있다:
[ 7, 92, 23, 9, -1, 0, 11, 6 ]
k-번째 작은 요소를 찾는 것이 더 쉽기 때문에, k = 4
로 설정하고 4번째로 작은 요소를 찾아보자.
배열을 먼저 정렬할 필요는 없다. 무작위로 피벗을 선택해보자. 예를 들어 11
을 선택하고, 이 피벗을 기준으로 배열을 분할한다. 다음과 같은 결과가 나올 수 있다:
[ 7, 9, -1, 0, 6, 11, 92, 23 ]
<------ 작은 값 큰 값 -->
보듯이, 11
보다 작은 모든 값은 왼쪽에, 큰 값은 오른쪽에 있다. 피벗 값 11
은 이제 최종 위치에 있다. 피벗의 인덱스는 5이므로, 4번째 작은 요소는 왼쪽 파티션 어딘가에 있다. 이제 나머지 배열은 무시할 수 있다:
[ 7, 9, -1, 0, 6, x, x, x ]
다시 무작위로 피벗을 선택해보자. 예를 들어 6
을 선택하고, 이 피벗을 기준으로 배열을 분할한다. 다음과 같은 결과가 나올 수 있다:
[ -1, 0, 6, 9, 7, x, x, x ]
피벗 6
은 인덱스 2에 위치하므로, 4번째 작은 요소는 오른쪽 파티션에 있어야 한다. 이제 왼쪽 파티션은 무시할 수 있다:
[ x, x, x, 9, 7, x, x, x ]
다시 무작위로 피벗을 선택해보자. 예를 들어 9
를 선택하고, 이 피벗을 기준으로 배열을 분할한다:
[ x, x, x, 7, 9, x, x, x ]
피벗 9
의 인덱스는 4이며, 이는 우리가 찾던 k와 정확히 일치한다. 작업이 완료되었다! 이 과정은 몇 단계만 걸렸고, 배열을 먼저 정렬할 필요가 없었다.
다음 함수는 이러한 아이디어를 구현한다:
public func randomizedSelect<T: Comparable>(_ array: [T], order k: Int) -> T {
var a = array
func randomPivot<T: Comparable>(_ a: inout [T], _ low: Int, _ high: Int) -> T {
let pivotIndex = random(min: low, max: high)
a.swapAt(pivotIndex, high)
return a[high]
}
func randomizedPartition<T: Comparable>(_ a: inout [T], _ low: Int, _ high: Int) -> Int {
let pivot = randomPivot(&a, low, high)
var i = low
for j in low..<high {
if a[j] <= pivot {
a.swapAt(i, j)
i += 1
}
}
a.swapAt(i, high)
return i
}
func randomizedSelect<T: Comparable>(_ a: inout [T], _ low: Int, _ high: Int, _ k: Int) -> T {
if low < high {
let p = randomizedPartition(&a, low, high)
if k == p {
return a[p]
} else if k < p {
return randomizedSelect(&a, low, p - 1, k)
} else {
return randomizedSelect(&a, p + 1, high, k)
}
} else {
return a[low]
}
}
precondition(a.count > 0)
return randomizedSelect(&a, 0, a.count - 1, k)
}
가독성을 위해 기능을 세 가지 내부 함수로 나누었다:
randomPivot()
은 무작위로 숫자를 선택해 현재 파티션의 끝에 위치시킨다(이는 Lomuto 분할 방식의 요구사항이다. 자세한 내용은 퀵 정렬 설명을 참고하라).
randomizedPartition()
은 퀵 정렬의 Lomuto 분할 방식을 사용한다. 이 함수가 완료되면, 무작위로 선택된 피벗은 배열에서 최종 정렬된 위치에 있게 된다. 이 함수는 피벗의 배열 인덱스를 반환한다.
randomizedSelect()
는 모든 주요 작업을 수행한다. 먼저 분할 함수를 호출한 후, 다음에 무엇을 할지 결정한다. 피벗의 인덱스가 찾고자 하는 k-번째 숫자와 같으면 작업이 완료된다. k
가 피벗 인덱스보다 작으면 왼쪽 파티션에서 재귀적으로 다시 시도한다. 마찬가지로 k-번째 숫자가 오른쪽 파티션에 있어야 할 경우도 처리한다.
꽤 멋지지 않은가? 일반적으로 퀵 정렬은 O(n log n) 알고리즘이지만, 배열의 점점 작은 부분만 분할하기 때문에 randomizedSelect()
의 실행 시간은 **O(n)**이 된다.
참고: 이 함수는 배열에서 k-번째 작은 요소를 계산한다. 여기서 k는 0부터 시작한다. k-번째 큰 요소를 원한다면
a.count - k
로 호출하면 된다.
Daniel Speiser가 작성. Matthijs Hollemans가 추가.