[백준] 2696 - 중앙값 구하기

최소 힙, 최대 힙 2가지를 사용하여 중앙값을 추출하는 문제이다.

풀이

풀이 과정의 경우 다음과 같이 풀었다.

1. t 만큼 repeat 하고, 중앙값을 구하기 위해 각 케이스마다 MaxHeap, MinHeap 구현
2. 시간 절약을 위해 StringTokenizer, StringBuilder 구현
3. MaxHeap 을 중앙값 이하로 유지하기 위해 MaxHeap 의 루트가 MinHeap 보다 큰 경우 두 개를 스왑
4. StringBuilder.append 가 10번이 되면 줄넘김

/*
* [입력]
* 홀 수 번째 읽는 경우 내용의 중앙 값을 출력.
* 첫째줄 - 테이스 케이스 개수 T
*
* 각 테스트 케이스 별
* 첫째 줄 - 수열의 크기 1 <= M(홀수) <= 9999
* 둘째 줄 - 수열의 원소 (한 줄에 최대 10개, 32비트 부호있는 정수)
*
* [출력]
* 첫째 줄 - 출력하는 중앙값의 개수
* 둘째 줄 - 홀수 번째 읽을 때마다 구한 중앙값 차례대로 공백으로 구분 출력 (최대 10줄)
*
* */
fun `2696-중앙값 구하기`(){
    val br = System.`in`.bufferedReader()
    val bw = System.out.bufferedWriter()
    val t = br.readLine().toInt()
    val stringBuilder = StringBuilder()

    repeat(t){ repeat ->
        val m = br.readLine().toInt()
        val maxHeap = PriorityQueue<Int>(compareByDescending { it })
        val minHeap = PriorityQueue<Int>()
        var appendedCount = 0
        var stringToken: StringTokenizer? = null

        if(repeat > 0) stringBuilder.appendLine()
        stringBuilder.appendLine("${m / 2 + 1}")

        for(i in 0 until m){
            if(i % 10 == 0) stringToken = StringTokenizer(br.readLine())
            if(stringToken == null) break

            val num = stringToken.nextToken().toInt()

            if(maxHeap.size == minHeap.size){
                maxHeap.offer(num)
            } else {
                minHeap.offer(num)
            }

            if(minHeap.isNotEmpty() && maxHeap.peek() > minHeap.peek()){
                val swapMax = maxHeap.poll()
                val swapMin = minHeap.poll()

                maxHeap.offer(swapMin)
                minHeap.offer(swapMax)
            }

            // 0 based Index 이므로 0이 남으면 홀수임
            if(i % 2 == 0){
                stringBuilder.append("${maxHeap.peek()} ")
                appendedCount++
            }

            // 추가가 10번 되었으면 한 라인 띄우기
            if(appendedCount == 10){
                stringBuilder.appendLine()
                appendedCount = 0
            }
        }
    }

    bw.write(stringBuilder.toString())
    bw.flush()
    bw.close()
    br.close()
}

후기

처음에는 힙을 이용하여 중앙값 구하는 방법에 대해 알 지 못해서 배열 + 파라메트릭 서치를 사용하려고 하였으나 아무리 생각해도 시간복잡도를 벗어났다.

그래서 이것저것 알아보니 최대 힙 + 최소 힙 두 가지를 사용하면 중간값을 알아낼 수 있다고 하더라.

다만 의문인건 힙을 사용하더라도 일반 CPU 컴퓨터를 사용하는 경우 T * M * O(log M) 이 되어 최악의 경우 1.4초 정도 가량이 나올텐데, 백준에서는 174ms 만 나왔다는 정도인데, 이게 맞나 싶어 다른곳도 찾아보니 다 동일하게 풀었더라.

앞으로 응용할 일이 많은 것 같은데 까먹지 않길 ...