BOJ-1202 보석 도둑 Swift

문제

세계적인 도둑 상덕이는 보석점을 털기로 결심했다.

상덕이가 털 보석점에는 보석이 총 N개 있다. 각 보석은 무게 Mi와 가격 Vi를 가지고 있다. 상덕이는 가방을 K개 가지고 있고, 각 가방에 담을 수 있는 최대 무게는 Ci이다. 가방에는 최대 한 개의 보석만 넣을 수 있다.

상덕이가 훔칠 수 있는 보석의 최대 가격을 구하는 프로그램을 작성하시오.

입력

첫째 줄에 N과 K가 주어진다. (1 ≤ N, K ≤ 300,000)

다음 N개 줄에는 각 보석의 정보 Mi와 Vi가 주어진다. (0 ≤ Mi, Vi ≤ 1,000,000)

다음 K개 줄에는 가방에 담을 수 있는 최대 무게 Ci가 주어진다. (1 ≤ Ci ≤ 100,000,000)

모든 숫자는 양의 정수이다.

출력

첫째 줄에 상덕이가 훔칠 수 있는 보석 가격의 합의 최댓값을 출력한다.

내가 푼 풀이

- 이번 문제는 많이 어려웠었다

- 보석과 가방의 개수가 최대 300,000으로 O(N^2)는 시간초과가 날 꺼라 생각했다.

- 가방에는 최대 한 개의 보석만 넣을 수 있다는 조건이 핵심이다.

- 가장 비싼 보석들만 골라서 해당 보석의 무게에 알맞은 가방 중 가장 최대무게가 적은 가방에 넣으면 된다.

- 오름차순으로 정렬된 가방에 넣는다면, 최대무게가 적은 가방순으로 가장 비싼 보석을 넣게 된다.

- 무게는 최대 10억이고, 가치는 최대 100만이므로 Uint64로 최댓값을 저장한다.

- Swift는 힙과 큐를 제공하지 않으므로 직접 구현해야 한다.

- 외워두도록 해야겠다.

- Heap, PriorityQueue는 구현을 위해 사용됐고, 아래에 주요 코드들이 있다.

 

• 가방의 최대무게가 m일 때, m보다 낮거나 같은 무게의 보석들을 우선순위큐에 저장 (보석의 가치 저장)
• 보석의 가치가 높은 순으로 저장된 우선순위 큐에서 가장 앞의 원소는 해당하는 가방에 넣을 수 있는 보석 중 가장 가치가 높은 보석을 말한다.
• 보석을 넣고, 우선순위큐에서 pop 후, 다음 가방을 쓴다.
• 최대무게보다 높은 보석들밖에 없다면 건너뛴다.

 

import Foundation

struct Heap<T: Comparable> {
    private var elements: [T] = []
    private let sortFunction: (T, T) -> Bool
    
    var isEmpty: Bool {
        return self.elements.count <= 1
    }
    
    var peek: T? {
        if self.isEmpty { return nil}
        return self.elements[1]
    }
    
    var count: Int {
        return self.elements.count - 1
    }
    
    init(elements: [T] = [], sortFunction: @escaping (T,T) -> Bool) {
        if !elements.isEmpty {
            self.elements = [elements.first!] + elements
        } else {
            self.elements = elements
        }
        self.sortFunction = sortFunction
        if elements.count > 1 {
            self.buildHeap()
        }
    }
    
    func leftChild(of index: Int) -> Int {
        return index * 2
    }
    func rightChild(of index: Int) -> Int {
        return index * 2 + 1
    }
    func parent(of index: Int) -> Int {
        return index / 2
    }
    
    mutating func swimUp(from index: Int) {
        var index = index
        while index != 1 && self.sortFunction(self.elements[index], self.elements[self.parent(of: index)]) {
            self.elements.swapAt(index, self.parent(of: index))
            index = self.parent(of: index)
        }
    }
    
    mutating func insert(node: T) {
        if self.elements.isEmpty {
            self.elements.append(node)
            return
        }
        self.elements.append(node)
        self.swimUp(from: self.elements.endIndex - 1)
    }
    
    mutating func diveDown(from index: Int) {
        var higherPriority = index
        let leftIndex = self.leftChild(of: index)
        let rightIndex = self.rightChild(of: index)
        
        if leftIndex < self.elements.endIndex && self.sortFunction(self.elements[leftIndex], self.elements[higherPriority]) {
            higherPriority = leftIndex
        }
        
        if rightIndex < self.elements.endIndex && self.sortFunction(self.elements[rightIndex], self.elements[higherPriority]) {
            higherPriority = rightIndex
        }
        
        if higherPriority != index {
            self.elements.swapAt(higherPriority, index)
            self.diveDown(from: higherPriority)
        }
    }
    
    mutating func buildHeap() {
        for index in (1...(self.elements.count / 2)).reversed() {
            self.diveDown(from: index)
        }
    }
    
    mutating func remove() -> T? {
        if self.elements.isEmpty { return nil}
        self.elements.swapAt(1, elements.endIndex - 1)
        let deleted = elements.removeLast()
        self.diveDown(from: 1)
        
        return deleted
    }
}

struct PriorityQueue<T: Comparable> {
    var heap: Heap<T>
    
    init(_ elements: [T] = [], _ sort: @escaping (T,T) -> Bool) {
        heap = Heap(elements: elements, sortFunction: sort)
    }
    
    var count: Int {
        return heap.count
    }
    
    var isEmpty: Bool {
        return heap.isEmpty
    }
    
    func top () -> T? {
        return heap.peek
    }
    
    mutating func clear() {
        while !heap.isEmpty {
            heap.remove()
        }
    }
    
    mutating func pop() -> T? {
        return heap.remove()
    }
    
    mutating func push(_ element: T) {
        heap.insert(node: element)
    }
}

var inputs = readLine()!.split(separator: " ").map{ Int($0)! }
var jewelys = [(m: Int, v: Int)]()
var bags = [Int]()
var result: UInt64 = 0

// 보석들을 튜플로 입력
for i in 0..<inputs[0] {
    var num = readLine()!.split(separator: " ").map{ Int($0)! }
    jewelys.append((num[0], num[1]))
}

// 가방 입력
for i in 0..<inputs[1] {
    bags.append(Int(readLine()!)!)
}

// 보석들의 무게를 오름차순으로 정렬
// 가방의 최대무게를 오름차순으로 정렬
jewelys.sort{ $0.0 < $1.0 }
bags.sort{ $0 < $1 }

// 우선순위큐 선언 (보석의 가치가 높은순)
// 같은 보석을 넣지 않도록 Idx를 따로 선언
// 우선순위큐의 구조상 isEmpty를 사용하기위해 따로 0을 push (필수가아님)
var idx = 0
var pq = PriorityQueue([Int](), { $0 > $1})
pq.push(0)

for i in 0..<bags.count {
    
    // 보석들중에 가방의 최대무게보다 작거나 같은 보석들을 큐에 저장
    while idx < jewelys.count {
        if jewelys[idx].m <= bags[i] {
            pq.push(jewelys[idx].v)     // 우선순위큐에 보석의 가치 push
        } else {
            break
        }
        idx += 1
    }
    // 해당 가방의 최대무게보다 작거나 같은 보석이 있을경우
    // 보석중의 가장 가치가높은(우선순위큐의 제일 첫번째 원소의값)을 결과값에 저장
    // 첫번쨰원소 pop
    if !pq.isEmpty {
        result += UInt64(pq.top()!)
        pq.pop()
        
    }
}

print(result)