k번째로 큰 요소 찾기 문제

정수 배열 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() 함수는 두 가지 매개변수를 받는다: 정수로 구성된 배열 ak이다. 이 함수는 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)
}

가독성을 위해 기능을 세 가지 내부 함수로 나누었다:

꽤 멋지지 않은가? 일반적으로 퀵 정렬은 O(n log n) 알고리즘이지만, 배열의 점점 작은 부분만 분할하기 때문에 randomizedSelect()의 실행 시간은 **O(n)**이 된다.

참고: 이 함수는 배열에서 k-번째 작은 요소를 계산한다. 여기서 k는 0부터 시작한다. k-번째 큰 요소를 원한다면 a.count - k로 호출하면 된다.

Daniel Speiser가 작성. Matthijs Hollemans가 추가.