[swift] Swift에서 정렬 함수를 사용하여 가장 큰 K개의 원소 찾기

05 Dec 2023

swift

Swift는 다양한 내장된 함수와 기능을 제공하여 배열을 정렬하고 가장 큰 K개의 원소를 찾을 수 있습니다. 이번 글에서는 Swift의 정렬 함수를 사용하여 가장 큰 K개의 원소를 찾는 방법에 대해 알아보겠습니다.

1. Swift 배열 정렬하기

Swift에서는 배열을 정렬하기 위해 sort() 함수를 사용할 수 있습니다. 이 함수는 기본적으로 배열을 오름차순으로 정렬합니다. 아래는 배열을 정렬하는 간단한 예제 코드입니다.

var numbers = [5, 2, 8, 1, 9]
numbers.sort()
print(numbers) // [1, 2, 5, 8, 9]

위의 코드에서는 numbers 배열을 sort() 함수를 사용하여 정렬하였습니다. 결과로 [1, 2, 5, 8, 9]가 출력됩니다.

2. 가장 큰 K개의 원소 찾기

가장 큰 K개의 원소를 찾기 위해서는 먼저 배열을 정렬한 후, 배열의 뒤에서 K개의 원소를 선택하면 됩니다. 아래는 이 과정을 단계별로 설명한 코드입니다.

var numbers = [5, 2, 8, 1, 9]
let k = 3

numbers.sort() // 배열을 오름차순으로 정렬
let largestK = Array(numbers.suffix(k)) // 배열의 뒤에서 K개의 원소 선택

print(largestK) // [5, 8, 9]

위의 코드에서는 numbers 배열을 정렬한 후, suffix(k) 함수를 사용하여 배열의 뒤에서 K개의 원소를 선택했습니다. 결과로 [5, 8, 9]가 출력됩니다.

3. 성능 개선을 위한 Heap 사용

만약 배열의 크기가 매우 크다면, 정렬 함수를 사용하여 배열 전체를 정렬하는 것은 성능상 이슈가 발생할 수 있습니다. 이때는 최소 힙(Min Heap)을 사용하여 성능을 개선할 수 있습니다.

struct Heap<Element: Comparable> {
    var elements: [Element] = []
    let sort: (Element, Element) -> Bool
    
    init(sort: @escaping (Element, Element) -> Bool) {
        self.sort = sort
    }
    
    var isEmpty: Bool {
        return elements.isEmpty
    }
    
    var count: Int {
        return elements.count
    }
    
    func leftChildIndex(ofParentAt index: Int) -> Int {
        return (2 * index) + 1
    }
    
    func rightChildIndex(ofParentAt index: Int) -> Int {
        return (2 * index) + 2
    }
    
    func parentIndex(ofChildAt index: Int) -> Int {
        return (index - 1) / 2
    }
    
    mutating func enqueue(_ element: Element) {
        var currentIndex = elements.count
        elements.append(element)
        
        while currentIndex > 0 {
            let parentIndex = self.parentIndex(ofChildAt: currentIndex)
            if sort(elements[currentIndex], elements[parentIndex]) {
                elements.swapAt(currentIndex, parentIndex)
                currentIndex = parentIndex
            } else {
                break
            }
        }
    }
    
    mutating func dequeue() -> Element? {
        guard !elements.isEmpty else {
            return nil
        }
        
        if elements.count == 1 {
            return elements.removeFirst()
        } else {
            let root = elements[0]
            elements[0] = elements.removeLast()
            var currentIndex = 0
            
            while true {
                let leftChildIndex = self.leftChildIndex(ofParentAt: currentIndex)
                let rightChildIndex = self.rightChildIndex(ofParentAt: currentIndex)
                
                var first = currentIndex
                if leftChildIndex < elements.count && sort(elements[leftChildIndex], elements[first]) {
                    first = leftChildIndex
                }
                
                if rightChildIndex < elements.count && sort(elements[rightChildIndex], elements[first]) {
                    first = rightChildIndex
                }
                
                if first == currentIndex {
                    break
                }
                
                elements.swapAt(currentIndex, first)
                currentIndex = first
            }
            return root
        }
    }
    
    func peek() -> Element? {
        return elements.first
    }
}

var numbers = [5, 2, 8, 1, 9]
let k = 3

var heap = Heap(sort: >) // 최대 힙(Max Heap)으로 사용하기 위해 '>'
for number in numbers {
    heap.enqueue(number)
    if heap.count > k {
        heap.dequeue()
    }
}

let largestK = heap.elements.reversed()
print(Array(largestK)) // [5, 8, 9]

위의 코드에서는 Heap 구조체를 정의하여 최대 힙으로 사용하고, 배열의 원소를 하나씩 확인하며 최대 힙에 추가합니다. 힙의 크기가 K를 초과하면 가장 작은 원소를 제거하고, 마지막에 힙의 원소를 역순으로 정렬하여 결과를 출력합니다.

이렇게 하면 배열의 크기와 상관없이 가장 큰 K개의 원소를 찾을 수 있으며, 성능도 개선됩니다.

요약

Swift에서는 정렬 함수를 사용하여 배열을 정렬할 수 있으며, 이를 활용하여 가장 큰 K개의 원소를 찾을 수 있습니다. 배열의 크기에 따라 성능을 개선하기 위해 최소 힙을 사용하는 방법도 있습니다. 이러한 기능을 활용하면 Swift에서 가장 큰 K개의 원소를 찾는 문제를 해결할 수 있습니다.