백준 - 보석도둑(Swift)

백준 - 보석도둑(Swift)

 

 

 문제 설명

 

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

1202번: 보석 도둑

 

 

 

 나의 풀이

 

가격을 기준으로 정렬하고 가방의 무게가 작은 것부터 탐색하는 아이디어는 올바른 접근이였지만 n과 k가 300,000까지이기 때문에 

for문안에 filter에서는 300,000*(300000log300000)의 시간복잡도를 갖게되기 때문에 맞지않았다.

 

 

import Foundation

let n = Int(readLine()!)!
let k = Int(readLine()!)!
var arr: [(Int, Int)] = []
var bags:[Int64] = []
var result = 0
for _ in 0..<n {
    let input = readLine()!.components(separatedBy: " ").map{Int($0)!}
    arr.append((input[0], input[1]))
}
arr.sort(by: {$0.1 == $1.1 ? $0.0 < $1.0 : $0.1 > $1.1 })

for _ in 0..<k {
    bags.append(Int(readLine()!)!)
}
bags.sort(by: <)

for bag in bags {
    let temp = arr.filter{ bag >= $0.0 }
    if temp.count > 0 {
        result += temp[0].1
        arr = arr.filter{ $0 != temp[0] }
    }
}
print(result)

 

그래서 다른사람의 풀이를 참고하여 다시 작성했다.

 

먼저 무게 별로 arr과 bags를 정렬하고 bags의 가방을 탐색하면서 조건에 부합하는 경우에 힙에 넣고 부합하지 않다면 while문을 탈출하는 식으로 두었다. 여기서 핵심은 idx가 전역변수로 되어있기 때문에 계속 초기화되는 것이 아니고 while문이 전에 끊긴 부분에서 다시 이어서 검사를 하게 된다. 

 

결국 처음했던 풀이는 가방마다 모든 arr을 돌아야 했기 때문에 O(가방의 개수 * 보석의 개수)였는데 우선순위 큐를 이용해서 한 풀이는 

O(가방의 개수 + 보석의 개수)가 되는 것이다. 

 

import Foundation

public struct Heap<T> {
    
    private var nodes = [T]()
    
    private var orderCriteria: (T, T) -> Bool
    
    public init(sort: @escaping (T, T) -> Bool) {
        // 최대 힙인지 최소 힙인지 기준을 잡는다.
        self.orderCriteria = sort
    }
    
    public init(array: [T], sort: @escaping (T, T) -> Bool) {
        self.orderCriteria = sort
        configureHeap(from: array)
    }
    
    public var isEmpty: Bool {
        return count == 0
    }
    
    public var count: Int {
        return nodes.count
    }
    
    public func peek() -> T? {
        return nodes.first
    }
    
    public mutating func insert(_ value: T) {
        nodes.append(value)
        shiftUp(nodes.count - 1)
    }
    
    public mutating func remove() -> T? {
        guard !nodes.isEmpty else { return nil }
        
        if nodes.count == 1 {
            return nodes.removeLast()
        } else {
            let value = nodes[0]
            nodes[0] = nodes.removeLast()
            shiftDown(0)
            return value
        }
    }
    
    public mutating func remove(at index: Int) -> T? {
        guard index < nodes.count else { return nil }
        
        let lastIndex = nodes.count-1
        if index != lastIndex {
            nodes.swapAt(index, lastIndex)
            shiftDown(from: index, until: lastIndex)
            shiftUp(index)
        }
        
        return nodes.removeLast()
    }
    
    // 변수를 직접 변경해야 하기 때문에 mutating 을 사용한다.
    private mutating func configureHeap(from array: [T]) {
        nodes = array
        
        /**
         * 힙 트리에서 n/2 - 1 은 하나 위 층위가 된다.
         * shiftDown을 하면서 자리를 잡을 것이기 때문에
         * 마지막 leaf 노드들은 할 필요가 없다.
         */
        for i in stride(from: nodes.count/2 - 1, through: 0, by: -1) {
            shiftDown(i)
        }
    }
    
    private func parentIndex(ofIndex i: Int) -> Int {
        return (i - 1) / 2
    }
    
    private func leftChildIndex(ofIndex i: Int) -> Int {
        return 2*i + 1
    }
    
    private func rightChildIndex(ofIndex i: Int) -> Int {
        return 2*i + 2
    }
    
    /**
     * shiftUp은 자신과 부모를 비교해 바꿔준다.
     */
    private mutating func shiftUp(_ index: Int) {
        var childIndex = index
        let child = nodes[childIndex] // 처음에 노드를 저장해두고 인덱스를 구한 후 바꿔준다.
        var parentIndex = self.parentIndex(ofIndex: index)
        
        while childIndex > 0 && orderCriteria(child, nodes[parentIndex]) {
            nodes[childIndex] = nodes[parentIndex]
            childIndex = parentIndex
            parentIndex = self.parentIndex(ofIndex: childIndex)
        }
        
        nodes[childIndex] = child
    }
    
    /**
     * shiftDown은 left, right 자식 중 적합한 녀석이 있으면 바꾸면서 바닥까지 내려간다.
     */
    private mutating func shiftDown(from index: Int, until endIndex: Int) {
        let leftChildIndex = self.leftChildIndex(ofIndex: index)
        let rightChildIndex = leftChildIndex + 1
        
        var first = index
        if leftChildIndex < endIndex && orderCriteria(nodes[leftChildIndex], nodes[first]) {
            first = leftChildIndex
        }
        if rightChildIndex < endIndex && orderCriteria(nodes[rightChildIndex], nodes[first]) {
            first = rightChildIndex
        }
        if first == index { return }
        
        nodes.swapAt(index, first)
        shiftDown(from: first, until: endIndex)
    }
    
    private mutating func shiftDown(_ index: Int) {
        shiftDown(from: index, until: nodes.count)
    }
    
}
let nk = readLine()!.split(separator: " ").map { Int(String($0))! }
var arr = [(m: Int, v: Int)]()
var bags = [Int]()
var result: Int = 0
for _ in 0..<nk[0] {
    let d = readLine()!.split(separator: " ").map { Int(String($0))! }
    arr.append((d[0], d[1]))
}
for _ in 0..<nk[1] {
    bags.append(Int(readLine()!)!)
}

arr.sort{ $0.m < $1.m }
bags.sort(by: <)

var heap = Heap<(Int, Int)>(array: [], sort: {$0.1 > $1.1})
var idx = 0
for bag in bags {
    while idx < arr.count, arr[idx].0 <= bag {
        heap.insert(arr[idx])
        idx += 1
    }
    let temp = heap.remove()
    result += temp?.1 ?? 0
}
print(result)

 

 

 피드백

 

항상 시간복잡도를 고려해서 문제를 풀어야한다. 이 문제같은 경우는 평범한 filter를 돌리기에는 너무 커서 우선순위 큐를 이용해야했다. 우선순위 큐는 삽입, 삭제가 logN이기 때문에 매우 작은 시간복잡도를 사용한다.

 

nil이 될 수 있는 변수같은 경우에는 ?를 붙여줘서 사용해야 한다. 

위의 예시로 let temp = heap.remove()같은 경우 temp는 옵셔널의 값이 되고 temp?.1을 사용함으로써 temp가 nil이 아니라면 언래핑해서 나오고 nil이라면 그 대체값으로 0을 리턴하도록 했다.