[백준 알고리즘/Kotlin] 1202번 보석 도둑 (feat. mutableListOf, PriorityQueue)

 

 

문제는 생각보다 심플한 듯하였습니다. 그래서 아래와 같은 로직을 세우고, 이대로 코드를 작성해야겠다고 생각했습니다.

1. 용량이 작은 가방부터 꺼낸다. (꺼냄과 동시에 자료형에서 제거해야 한다.)
2. 꺼낸 가방의 용량보다 작거나 같은 보석 중 가장 가격이 높은 보석을 꺼낸다. (꺼냄과 동시에 자료형에서 제거해야 한다.)
3. 해당하는 보석의 가격을 확인하여 result에 + 한다.
4. 더이상더 이상 꺼낼 가방이 없을 때 종료 (혹은 더 이상 꺼낼 주얼리가 없을 때 종료)

 

1. 시간 초과된 코드

import java.io.BufferedReader
import java.io.InputStreamReader
import java.util.PriorityQueue


fun main() {
    val br = BufferedReader(InputStreamReader(System.`in`))

    val (n, k) = br.readLine().split(" ").map { it.toInt() }
    
    // 가격 높은 순으로 보석 저장
    val jewels = PriorityQueue<Pair<Int, Int>>(compareByDescending { it.second })
    repeat(n) {
        val (weight, price) = br.readLine().split(" ").map { it.toInt() }
        jewels.add(weight to price)
    }
    
    // 용량 작은 순으로 가방 저장
    val bags = PriorityQueue<Int>()
    repeat(k) {
        val capacity = br.readLine().toInt()
        bags.add(capacity)
    }
    
    var result = 0L
    
    while (bags.isNotEmpty() && jewels.isNotEmpty()) {
        val bagCapacity = bags.poll() // 현재 가방의 용량
        val copyOfJewels = PriorityQueue(jewels) // 원본 훼손하지 않기 위한 복사본
        var eligibleJewel: Pair<Int, Int>? = null // 조건에 맞는 보석
        
        while (copyOfJewels.isNotEmpty()) { // 보석 무게 ≤ 가방 용량 중 가장 가격이 높은 보석
            val jewel = copyOfJewels.poll()
            
            if (jewel.first <= bagCapacity) {
                eligibleJewel = jewel
                jewels.remove(jewel) // 원본에서 해당 보석 삭제
                break
            }
        }
        
        if (eligibleJewel != null) {
            result += eligibleJewel.second
        }
    }

    println(result)
    br.close()
}

 

왜 시간 초과가 났는지 코드를 확인해보도록 하죠.

 

1-1. 우선순위 큐를 사용한 이유

보석과 가방을 우선순위 큐(PriorityQueue)에 저장한 이유는

(1) 특정한 조건을 기준으로 정렬하고 (2) 꺼냄과 동시에 제거해야 하는 poll() 메서드를 사용하기 위해서입니다.

val (n, k) = br.readLine().split(" ").map { it.toInt() }

// 가격 높은 순으로 보석 저장
val jewels = PriorityQueue<Pair<Int, Int>>(compareByDescending { it.second })
repeat(n) {
    val (weight, price) = br.readLine().split(" ").map { it.toInt() }
    jewels.add(weight to price)
}

// 용량 작은 순으로 가방 저장
val bags = PriorityQueue<Int>()
repeat(k) {
    val capacity = br.readLine().toInt()
    bags.add(capacity)
}

 

1-2. while 문의 종료 조건

더 이상 꺼낼 가방이 없을 때,혹은 더이상 꺼낼 주얼리가 없을 때 종료해야 하기 때문에 아래와 같은 조건을 세웠습니다.

while (bags.isNotEmpty() && jewels.isNotEmpty()) {
    ...
}

 

1-3. 주요 로직 체크하기

var result = 0L

while (bags.isNotEmpty() && jewels.isNotEmpty()) {
    val bagCapacity = bags.poll() // 현재 가방의 용량
    val copyOfJewels = PriorityQueue(jewels) // 원본 훼손하지 않기 위한 복사본
    var eligibleJewel: Pair<Int, Int>? = null // 조건에 맞는 보석

    while (copyOfJewels.isNotEmpty()) { // 보석 무게 ≤ 가방 용량 중 가장 가격이 높은 보석
        val jewel = copyOfJewels.poll()

        if (jewel.first <= bagCapacity) {
            eligibleJewel = jewel
            jewels.remove(jewel) // 원본에서 해당 보석 삭제
            break
        }
    }

    if (eligibleJewel != null) {
        result += eligibleJewel.second
    }
}

 

로직과 함께 코드를 살펴보죠.

1. 용량이 적은 가방부터 꺼낸다. (꺼냄과 동시에 자료형에서 제거해야 한다.)
2. 꺼낸 가방의 용량보다 작거나 같은 보석 중 가장 가격이 높은 보석을 꺼낸다. (꺼냄과 동시에 자료형에서 제거해야 한다.)
3. 해당하는 보석의 가격을 확인하여 result에 + 한다.
4. 더 이상 꺼낼 가방이 없을 때 종료 (혹은 더이상 꺼낼 주얼리가 없을 때 종료)

 

① 용량이 적은 순으로 정렬되어 있는 bags에서 poll()을 통해 꺼냄과 동시에 bags에서 제거합니다.

val bagCapacity = bags.poll()

 

② copyOfJewels 복사본을 만듭니다.

jewels(원본)에서는 꺼낸 가방의 용량보다 작거나 같은 보석을 찾아 해당 보석하는 것만 삭제해야 합니다.

그러나 각 보석의 무게를 꺼내어 가방의 용량보다 작거나 같은지 확인하기 위해서는 원본을 훼손하기 않기 위해 복사본이 필요합니다.

val copyOfJewels = PriorityQueue(jewels)

 

③ 가방의 용량보다 작거나 같은 보석 중 가장 가격이 높은 적합한 보석을 찾기 위한 변수를 선언해 둡니다.

var eligibleJewel: Pair<Int, Int>? = null

 

④ 꺼낸 가방의 용량보다 작거나 같은 보석 중 가장 가격이 높은 보석을 꺼냅니다.

이 조건을 만족시키기는 보석을 찾는다면 곧바로 종료합니다.

while (copyOfJewels.isNotEmpty()) {
    val jewel = copyOfJewels.poll()

    if (jewel.first <= bagCapacity) {
        eligibleJewel = jewel
        jewels.remove(jewel)
        break
    }
}

 

⑤ 만약 보석을 찾았다면 해당 보석의 second 값인 가격(price)을 result에 더합니다.

if (eligibleJewel != null) {
    result += eligibleJewel.second
}

 

그러나 시간 초과가 나는 이유는 로직이 잘못되었거나, 더 나은 방법이 있다는 의미겠죠.

수정한 최종 코드를 보도록 하겠습니다.

 

2. 최종 코드

import java.io.BufferedReader
import java.io.InputStreamReader
import java.util.PriorityQueue

data class Jewel(val weight: Int, val price: Int) : Comparable<Jewel> {
    override fun compareTo(other: Jewel) = other.price.compareTo(price)
}

fun main() {
    val br = BufferedReader(InputStreamReader(System.`in`))

    val (n, k) = br.readLine().split(" ").map { it.toInt() }
    val jewels = mutableListOf<Jewel>()
    val bags = mutableListOf<Int>()
    
    repeat(n) {
        val (weight, price) = br.readLine().split(" ").map { it.toInt() }
        jewels.add(Jewel(weight, price))
    }
    
    repeat(k) {
        val capacity = br.readLine().toInt()
        bags.add(capacity)
    }
    
    jewels.sortBy { it.weight } // 무게 작은 순으로 보석 저장
    bags.sort() // 용량 작은 순으로 가방 저장
	
    val jewelsPriorityQueue = PriorityQueue<Jewel>()
    var jewelsIdx = 0
    var result: Long = 0
    
    for (i in 0 until k) {
        while (jewelsIdx < n && jewels[jewelsIdx].weight <= bags[i]) {
            jewelsPriorityQueue.add(jewels[jewelsIdx++])
        }
        if (jewelsPriorityQueue.isNotEmpty()) {
            result += jewelsPriorityQueue.poll().price
        }
    }
    
    println(result)
    
    br.close()
}

 

2-1. 데이터 모델 추가

코틀린에는 데이터 모델을 간단히 만들 수 있다는 이점이 있죠.

Pair의 first, second가 아닌 Jewel의 weight, price로 코드의 가독성과 함께 보석의 의미를 명확하게 만들도록 하겠습니다.

그리고 Comparable을 구현하여 Jewel는 price(가격)을 기준으로 내림차순 정렬하도록 하겠습니다.

data class Jewel(val weight: Int, val price: Int) : Comparable<Jewel> {
    override fun compareTo(other: Jewel) = other.price.compareTo(price)
}

 

2-2. 데이터 저장 구조 변경

PriorityQueue → mutableListOf

val jewels = mutableListOf<Jewel>()
val bags = mutableListOf<Int>()

repeat(n) {
    val (weight, price) = br.readLine().split(" ").map { it.toInt() }
    jewels.add(Jewel(weight, price))
}

repeat(k) {
    val capacity = br.readLine().toInt()
    bags.add(capacity)
}

jewels.sortBy { it.weight } // 무게 작은 순으로 보석 저장
bags.sort() // 용량 작은 순으로 가방 저장

1. 시간 복잡도
PriorityQueue의 삽입(add 또는 offer)과 삭제(poll) 연산의 시간 복잡도는 O(log n)
mutableListOf의 원소의 삽입과 삭제 연산의 시간 복잡도는 O(1)

2. 용도
PriorityQueue는 삽입, 삭제에 따라 자동으로 정렬된 순서가 필요할 때 유용하지만,
문제에서는 새롭게 순서를 정렬하는 것이 아니라 순서를 유지하며 앞에 있는 원소를 삭제하면 되기 때문에
mutableListOf가 더 적합하다고 판단하였습니다.

3. 중복 원소 여부
PriorityQueue는 중복을 허용하지 않고, mutableListOf는 중복을 허용하는데요.
문제에 어디에도 중복되지 않는 가방의 용량이 나온다는 말이 없으므로 (즉, 중복될 수 있으므로)
mutableListOf를 사용하는 게 좋겠습니다.

4. jewels는 왜 무게 작은 순으로 저장하는 것으로 바뀌었나요? (변경된 로직)
가방과의 매칭 작업을 최적화하기 위함입니다.
보석을 무게순으로 정렬하면, (작은 용량 순으로 저장된) 현재 가방의 용량과 일치하는 보석을 찾는 과정이
효율적으로 이뤄지기 때문에, 로직을 변경하였습니다.

 

2-3. 주요 로직 체크하기

val jewelsPriorityQueue = PriorityQueue<Jewel>()
var jewelsIdx = 0
var result: Long = 0

for (i in 0 until k) {
    while (jewelsIdx < n && jewels[jewelsIdx].weight <= bags[i]) {
        jewelsPriorityQueue.add(jewels[jewelsIdx++])
    }
    if (jewelsPriorityQueue.isNotEmpty()) {
        result += jewelsPriorityQueue.poll().price
    }
}

println(result)

jewelsPriorityQueue
로직을 위해 현재 가방의 용량보다 작거나 같은 모든 보석을 저장하는 용도로 사용됩니다.
만약 jewelsPriorityQueue가 비어있지 않다면, 우선순위 큐에서 최고 가격의 보석을 꺼내어 result에 더합니다.
이러한 로직으로 사용해 시간을 최적화시킬 수 있었습니다.

jewelsIdx
보석의 인덱스를 통해 보석을 추가하거나 접근하는 것을 더 명시적으로 표현하였습니다.

 

마무리

우선순위 큐의 올바른 사용처를 알게 된 좋은 문제였는데요. 조금 더 코드에 대해 고민하는 습관을 가져야겠습니다 :)