N번째 큰 수(우선순위 힙, FileIO)[백준 실버 3]

 

시간복잡도에 가장 큰 영향을 주는 N값의 범위가 1<N<1500임을 확인하여, 이중 for문을 통해 1차원 배열에 모든 숫자들을 append 처리한 후, sort()를 실행해 5번째로 큰 숫자를 출력하여도 문제없을 것이라 생각하고 시도해보았습니다.

let N = Int(readLine()!)!
var arr = [Int]()
for _ in 0..<N {
	let input = readLine()!.split(separator: " ").map{Int($0)!}
    for id in 0..<N {
    	arr.append(input[id])
    }
}
arr.sort(by:>)
print(arr[4])

하지만, 위 코드로 성공했으면 실버 3이 아니겠죠.

 

위 코드의 시간복잡도는 아래와같이 분석해볼 수 있습니다.

 

1. 입력 처리 부분: 이중 For문으로 append() 실행

O(N^2)

 

2. 정렬 부분 : sort()

여기서 sort 메서드느 O( n log n )를 가지며, arr 크기는 N^2입니다

O(N^2 * logN^2) = O(N^2 * 2logN) = O(N^2 * logN)

 

따라서 전체 시간복잡도는 아래와 같습니다.

O(N²) + O(N² log N) + O(1) = O(N² log N)

 

그다음으로 시도해본 것은 우선순위 큐를 활용하는 것이였습니다.

우선순위 큐는 일반적인 FIFO 방식의 Queue와 달리, 각 요소가 우선순위를 가지고 있어, 양 끝 우선순위를 가진 요소가 먼저 처리되는 자료구조입니다.

일반적으로 완전 이진 트리 형태를 지닌 힙을 통해 구현되며, 이때 주요 연산의 시간복잡도는 아래와 같습니다.

삽입(insert): O(log n)

삭제(remove): O(log n)

최대/최소값 확인: O(1)

 

따라서 데이터의 특정 순서나 우선순위가 중요하고, 전체 데이터의 정렬이 필요하지 않은 상황에서 매우 효율적인 자료구조입니다.

 

해당 알고리즘 문제는 n번째까지만 정렬이 필요하며 그 이후부터는 정렬이 불필요하기에, 우선순위 큐로 성능 이점을 챙길 수 있다고 판단했습니다. 따라서, 아래와 같이 Heap을 활용한 우선순위 큐를 구현해보았습니다.

struct Heap<T: Comparable> {
    private(set) var items: [T] // 정렬되는 요소 배열
    private let sortFunc: (T, T) -> Bool // 정렬 기준 클로저
    
    /// 정렬 함수만 제공하여 빈 힙 생성할 경우
    /// 내부적으로 클로저를 저장하거나 비동기적으로 실행할 가능성이 있기 때문에, 함수 실행 종료 이후에도 실행될 수 있도록 escaping 명시
    init(sortFunc: @escaping (T, T) -> Bool) {
        self.items = []
        self.sortFunc = sortFunc
    }
    
    /// 배열과 정렬 함수를 함께 제공해 힙 생성할 경우
    init(array: [T], sortFunc: @escaping (T, T) -> Bool) {
        self.items = array
        self.sortFunc = sortFunc
        // 마지막 노드 바로 위 노드들부터 루트(0번째)까지 거꾸로 각 노드에 대해
        // 부모노드가 주체가 되어 자식들과 비교하면서 아래로 내려감
        for i in stride(from: items.count / 2 - 1, through: 0, byㄹ: -1) { 
        	siftDown(from: i) 
        }
    }
    
    /// 새로운 요소 추가
    mutating func insert(_ element: T) {
        items.append(element)
        swimUp(from: items.count - 1)
    }
    
    /// 루트 요소를 제거하고 반환
    mutating func remove() -> T? {
        guard !items.isEmpty else { return nil }
        
        items.swapAt(0, items.count - 1) // 마지막 요소랑 바꾸고 제거
        let element = items.removeLast()
        if !items.isEmpty { siftDown(from: 0) }
        
        return element
    }
    
    /// 자식 노드가 주체가 되어 부모와 비교하면서 위로 올라감
    private mutating func swimUp(from index: Int) { // 요소를 위로 이동
        var child = index
        var parent = parentIndex(of: child)
        
        while child > 0 && sortFunc(items[child], items[parent]) {
            items.swapAt(child, parent)
            child = parent
            parent = parentIndex(of: child)
        }
    }
    
    /// 요소를 아래로 이동
    private mutating func siftDown(from index: Int) { // 요소를 아래로 이동
        var parent = index
        while true {
            let leftChild = leftChildIndex(of: parent)
            let rightChild = rightChildIndex(of: parent)
            var candidate = parent
            
            if leftChild < items.count && sortFunc(items[leftChild], items[candidate]) { candidate = leftChild }
            if rightChild < items.count && sortFunc(items[rightChild], items[candidate]) { candidate = rightChild }
            if candidate == parent { return }
            
            items.swapAt(parent, candidate)
            parent = candidate
        }
    }
    
    private func parentIndex(of index: Int) -> Int { return (index - 1) / 2 }
    private func leftChildIndex(of index: Int) -> Int { return 2 * index + 1 }
    private func rightChildIndex(of index: Int) -> Int { return 2 * index + 2 }
}

 

우선순위 큐 구조체를 생성하고 난 이후, 아래와 같이 코드를 구성했습니다,

let N = Int(readLine()!)!
var heap = Heap<Int>(sortFunc: <)
for _ in 0..<N {
    let nums = readLine()!.split(separator: " ").map{Int($0)!}
    for num in nums {
        if heap.items.count<N {
            heap.insert(num)
        } else {
            if num > heap.items.first! {
                heap.remove()
                heap.insert(num)
            }
        }
    }
}

print(heap.items.first!)

모든 숫자 N*N개를 순회하며, 우선순위 큐에 요소를 삽입하는 것을 반복하는 구조를 구성하였는데요.

여기서 우선순위 큐의 최대 길이가 N을 넘기지않도록 하였고, 최소우선순위 큐로 정렬기준을 설정해주었습니다.

 

이로 인해, 숫자를 순회할때마다 전체 숫자들 중 최대값이 될 가능성이 있는 숫자들만 힙 내부 더 작은 숫자와 대체되게끔 알고리즘을 설계해, 불필요한 우선순위 큐 연산을 최소화하고자 했습니다.

마지막으로, 힙의 루트에 위치한 요소를 출력하도록 해, N번째로 큰 숫자가 출력되도록 했습니다.

앞서 말씀드렸듯이 힙의 모든 연산은 logN의 시간복잡도를 지닙니다. 여기서 N은 힙의 현재 크기를 의미하죠.

이중 for문으로 인해 N^2번 순회하며 힙을 접근하고 있는만큼 힙 자료구조 연산은 logN * N^2, 정확히는 log1 + log2 + log3 + log4+  log5 + ,,,가 될 것입니다.

 

따라서, 위 코드의 시간복잡도는 아래와 같습니다.

N * N * logN = N^2logN

앞서 말씀드렸던 일반 정렬 알고리즘의 시간복잡도와 표면상 동일해보이지만, 메모리 사용량 측면에서는 O(N)으로 효율적입니다.

즉, 일반정렬 알고리즘의 경우, 모든 N^2개의 숫자를 배열에 저장하고 전체를 정렬하는 반면, 힙 코드는 최대 N개만 메모리에 저장하고 부분정렬합니다.

또한, N^2개의 모든 숫자에 대해 정렬을 수행했던 일반정렬때와 달리, 힙을 사용할때에는 N번째로 큰 수보다 작은 숫자들은 연산하지 않고 스킵했기에 절대적인 연산량 또한 적죠.

 

하지만 가차없이 실패!

 

이후, 구글링하여 해당 문제를 접근한 다른분의 코드를 참고해보았는데요.

https://gist.github.com/JCSooHwanCho/30be4b669321e7a135b84a1e9b075f88

ps할 때 입력을 한꺼번에 받기 위한 유틸리티 클래스. fread의 swift 버전.

final class FileIO {
    private let buffer:[UInt8] // 파일의 전체 내용을 바이트 배열로 저장할 변수
    private var index: Int = 0

    // 파일의 모든 내용을 읽어 buffer에 저장하고, 범위 체크를 위해 끝에 0 추가
    init(fileHandle: FileHandle = FileHandle.standardInput) {
        buffer = Array(try! fileHandle.readToEnd()!)+[UInt8(0)] // 인덱스 범위 넘어가는 것 방지
    }
    
    // 한 바이트를 읽는 기본 함수 : inline 최적화로 함수 호출 비용을 줄임
    @inline(__always) private func read() -> UInt8 {
        defer { index += 1 } // defer로 함수 종료 시 인덱스를 증가
        return buffer[index]
    }
    
    // 정수를 읽는 함수의 시작부분
    @inline(__always) func readInt() -> Int {
        var sum = 0
        var now = read() // now: 현재 읽은 바이트
        var isPositive = true // 양수/음수 판별 플래그

        while now == 10 || now == 32 { now = read() } // 공백과 줄바꿈 무시
        if now == 45 { isPositive.toggle(); now = read() } // 음수 처리
        // ASCII 코드 48('0')~57('9')의 범위에서 실제 숫자로 변환
        while now >= 48, now <= 57 {
            sum = sum * 10 + Int(now-48)
            now = read()
        }

        return sum * (isPositive ? 1:-1)
    }

    @inline(__always) func readString() -> String {
        var now = read()

        while now == 10 || now == 32 { now = read() } // 공백과 줄바꿈 무시
        let beginIndex = index-1

        while now != 10, now != 32, now != 0 { now = read() }
        // 다음 공백/줄바꿈/끝(0)이 나올 때까지 읽습니다.
        return String(bytes: Array(buffer[beginIndex..<(index-1)]), encoding: .ascii)!
    }

    @inline(__always) func readByteSequenceWithoutSpaceAndLineFeed() -> [UInt8] {
        var now = read()

        while now == 10 || now == 32 { now = read() } // 공백과 줄바꿈 무시
        let beginIndex = index-1

        while now != 10,
              now != 32,
              now != 0 { now = read() }

        return Array(buffer[beginIndex..<(index-1)])
    }
}

 

위 코드는 파일의 내용을 바이트배열로 한번에 받아들인 후, 바이트 단위로 직접 처리하는 클래스입니다. 처리하는 도중에 문자열 변환 및 타입 변환으로 인한 오버헤드가 적으며, 각 줄마다 새로운 문자열을 생성하는 readLine()과 달리 FileIO는 전체 입력을 하나의 바이트 배열로 관리하여 메모리 할당을 최소화합니다. 마지막으로, FileOP는 한번의 시스템 콜로 모든 데이터를 읽어들여, 대량의 입력이 있는 경우 훨씬 빠른 성능을 보인다고 합니다.

 

최종적으로 FileIO를 통해 입력을 받아들이고, Heap자료구조를 통해 정렬해 아래와 같이 코드를 작성할 수 있었습니다.

import Foundation

struct Heap<T: Comparable> {
    private(set) var items: [T]
    private let sortFunc: (T, T) -> Bool
    
    /// Heap<Int> { $0 > $1 } 빈 힙 생성 // 정렬 함수만 제공하여 빈 힙 생성
    ///- Parameter sortFunc 내부적으로 클로저를 저장하거나 비동기적으로 실행할 가능성이 있기 때문에, 함수 실행 종료 이후에도 실행될 수 있기에 escaping 명시
    init(sortFunc: @escaping (T, T) -> Bool) {
        self.items = []
        self.sortFunc = sortFunc
    }
    
    /// let arrayHeap = Heap(array: [7, 3, 9, 1, 4]) { $0 > $1 } // 배열과 정렬 함수를 함께 제공하여 힙 생성
    init(array: [T], sortFunc: @escaping (T, T) -> Bool) {
        self.items = array
        self.sortFunc = sortFunc
        
        for i in stride(from: items.count / 2 - 1, through: 0, by: -1) { siftDown(from: i) }
    }
    
    /// 새로운 요소 추가
    mutating func insert(_ element: T) {
        items.append(element)
        swimUp(from: items.count - 1)
    }
    
    /// 루트 요소를 제거하고 반환
    mutating func remove() -> T? {
        guard !items.isEmpty else { return nil }
        
        items.swapAt(0, items.count - 1) // 마지막 요소랑 바꾸고 제거
        let element = items.removeLast()
        if !items.isEmpty { siftDown(from: 0) }
        
        return element
    }
    
    /// 요소를 위로 이동
    private mutating func swimUp(from index: Int) { // 요소를 위로 이동
        var child = index
        var parent = parentIndex(of: child)
        
        while child > 0 && sortFunc(items[child], items[parent]) {
            items.swapAt(child, parent)
            child = parent
            parent = parentIndex(of: child)
        }
    }
    
    /// 요소를 아래로 이동
    private mutating func siftDown(from index: Int) { // 요소를 아래로 이동
        var parent = index
        while true {
            let leftChild = leftChildIndex(of: parent)
            let rightChild = rightChildIndex(of: parent)
            var candidate = parent
            
            if leftChild < items.count && sortFunc(items[leftChild], items[candidate]) { candidate = leftChild }
            if rightChild < items.count && sortFunc(items[rightChild], items[candidate]) { candidate = rightChild }
            if candidate == parent { return }
            
            items.swapAt(parent, candidate)
            parent = candidate
        }
    }
    
    private func parentIndex(of index: Int) -> Int { return (index - 1) / 2 }
    private func leftChildIndex(of index: Int) -> Int { return 2 * index + 1 }
    private func rightChildIndex(of index: Int) -> Int { return 2 * index + 2 }
}

final class FileIO {
    private let buffer:[UInt8]
    private var index: Int = 0

    init(fileHandle: FileHandle = FileHandle.standardInput) {
        buffer = Array(try! fileHandle.readToEnd()!)+[UInt8(0)] // 인덱스 범위 넘어가는 것 방지
    }

    @inline(__always) private func read() -> UInt8 {
        defer { index += 1 }

        return buffer[index]
    }

    @inline(__always) func readInt() -> Int {
        var sum = 0
        var now = read()
        var isPositive = true

        while now == 10
                || now == 32 { now = read() } // 공백과 줄바꿈 무시
        if now == 45 { isPositive.toggle(); now = read() } // 음수 처리
        while now >= 48, now <= 57 {
            sum = sum * 10 + Int(now-48)
            now = read()
        }

        return sum * (isPositive ? 1:-1)
    }

    @inline(__always) func readString() -> String {
        var now = read()

        while now == 10 || now == 32 { now = read() } // 공백과 줄바꿈 무시
        let beginIndex = index-1

        while now != 10,
              now != 32,
              now != 0 { now = read() }

        return String(bytes: Array(buffer[beginIndex..<(index-1)]), encoding: .ascii)!
    }

    @inline(__always) func readByteSequenceWithoutSpaceAndLineFeed() -> [UInt8] {
        var now = read()

        while now == 10 || now == 32 { now = read() } // 공백과 줄바꿈 무시
        let beginIndex = index-1

        while now != 10,
              now != 32,
              now != 0 { now = read() }

        return Array(buffer[beginIndex..<(index-1)])
    }
}

let fIO = FileIO()
let N = fIO.readInt()
var heap = Heap<Int>(sortFunc: <)

for _ in 0..<N*N {
    let num = fIO.readInt()
    if heap.items.count<N {
        heap.insert(num)
    } else {
        if num > heap.items.first! {
            heap.remove()
            heap.insert(num)
        }
    }
}

print(heap.items[0])

비록 제 힘으로는 풀지못한 반쪽짜리 결과물이지만, Swift 기준 백준 랭킹 2위네요!! 먼가 의미없는 순위긴하지만 은근 기분이 좋습니다. 그나저나, 다들 코드 길이를 보니 고생을 하신것 같습니다...

 

감사합니다!