만득

코딩 공부 정리용

​

© 2021. All rights reserved.

백준 1655 가운데를 말해요

23 Jun 2021

https://www.acmicpc.net/problem/1655

입력이 하나씩 들어올때마다 여지껏 들어온 입력들 중에서 중간값을 출력해야하는 문제이다. 만약 짝수개라면 두 수 중에서 작은 수를 말해야 한다.

생각보다 아이디어를 생각하기 어려운 문제였다. 두 개의 힙을 준비해서 정 중앙 값을 기준으로 절반은 왼쪽에 절반은 오른쪽 힙에 집어넣는다고 하자. 만약 왼쪽 heap이 max heap이고 오른쪽 heap이 min heap이라면 두 heap은 항상 정중앙에 가까운 값을 root에 위치시킬 것이다. 홀수번째에는 max heap에 짝수번째에는 min heap에 값을 집어넣는다고 하자.

만약 정상적으로 코드가 동작해서 중간값을 기준으로 작은 값들은 모두 max heap에 들어있고 큰 값들은 모두 min heap에 들어있다고 하자. 그렇다면 새로운 입력이 들어왔을 때 그 값이 max heap에 들어갔다고 하자. 그런데 새로운 입력이 중간값보다 훨씬 커서 min heap에 위치해야한다면 이를 옮겨줘야 할 것이다. 당연히 새로운 입력은 max heap의 모든 값보다 크기 때문에 max heap의 root에 위치할 것이다. 따라서 과정은 다음과 같다.

1. max heap의 root와 min heap의 root를 바꾼다.
3. 기존의 min heap의 최솟값은 max heap의 새로운 최댓값이 될 것이고 root에 위치시키면 된다.
4. min heap의 새로운 root에 shift down 함수를 적용해서 적절한 위치를 찾아준다.

반대로 새로운 입력이 min heap에 들어갔지만 중간값보다 작아서 옮겨줘야 할 경우에도 위와 동일하게 적용을 시켜준다. 다만 이번에는 max heap에 shift down을 적용해서 적절한 위치를 찾아준다. 최종적으로 중간값은 max heap의 root에 위치할 것이다.

struct Heap<T> {
    var arr:[T]
    var compare:(T, T) -> Bool
    var root:T? {
        get {
            if arr.isEmpty {
                return nil
            }
            else {
                return arr[0]
            }
        }
    }
    mutating func shiftUp(a: Int) {
        var index = a
        while index > 0 {
            let parent = (index-1)/2
            if compare(arr[index], arr[parent]) {
                arr.swapAt(index, parent)
                index = parent
            }
            else {
                break
            }
        }
    }
    mutating func shiftDown(a: Int) {
        var index = a
        var child = 2*a+1
        let count = arr.count
        while child < count {
            if child + 1 < count {
                child = compare(arr[child], arr[child+1]) ? child : child+1
            }
            if compare(arr[child], arr[index]) {
                arr.swapAt(index, child)
                index = child
                child = 2*index+1
            }
            else {
                break
            }
        }
    }
    mutating func insert(_ a: T) {
        arr.append(a)
        shiftUp(a: arr.count-1)
    }
    mutating func pop() -> T? {
        if arr.isEmpty {
            return nil
        }
        else {
            arr.swapAt(0, arr.count-1)
            let result = arr.removeLast()
            shiftDown(a: 0)
            return result
        }
    }
    init(arr: [T] = [], compare: @escaping (T, T) -> Bool) {
        self.arr = arr
        self.compare = compare
    }
}

var minHeap = Heap<Int>{$0<$1}
var maxHeap = Heap<Int>{$0>$1}

let n = Int(readLine()!)!
if n == 1 {
    print(Int(readLine()!)!)
}
else {
    var str = ""
    maxHeap.insert(Int(readLine()!)!)
    str += "\(maxHeap.root!)\n"
    for i in 2...n {
        let k = Int(readLine()!)!
        if i%2 == 0 {minHeap.insert(k)}
        else {maxHeap.insert(k)}
        let a = maxHeap.root!
        let b = minHeap.root!
        if a > b {
            minHeap.arr[0] = a
            maxHeap.arr[0] = b
        }
        if i%2 == 0 {maxHeap.shiftDown(a: 0)}
        else {minHeap.shiftDown(a: 0)}
        str += "\(maxHeap.root!)\n"
    }
    print(str)
}

Related posts