[백준 2696번] [Kotlin] 중앙값 구하기

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

2696번: 중앙값 구하기

난이도 : 골드 2
태그 : 자료구조, 우선순위 큐

 

 

설명

수를 차례로 입력받아가면서,

홀수 번째 수를 입력받았을 때 중앙값을 출력하는 문제입니다.

 

음.. 매번 정렬을 하면 당연히 시간초과가 날 것 같고,

정렬을 하는게 목적이 아니라, 단순히 중앙값을 찾는게 목적이기 때문에

우선순위 큐를 두 개 사용하면 해결할 수 있을 것 같습니다.

 

저는 일단, 절반을 기준으로 작은 값을 담는 lowQ, 큰 값을 담는 highQ 두 개를 만들어봤습니다.

lowQ는 내림차순, highQ는 오름차순 정렬로 지정하고,

 

이 두개는 지금까지 입력받은 숫자가 짝수개일 때는 크기가 같아야 하고,

홀수일 경우에는 lowQ가 highQ 보다 1 만큼 사이즈가 더 크도록 만든 뒤,

홀수개일 때 마다 lowQ의 맨 윗 숫자를 출력해주는 방식으로 작동합니다.

 

1 5 6 9 2 이 차례로 들어온다고 해봅시다.

 

1회차) 입력받은 수 - 1 , 남은 수 - 5 6 9 2

lowQ - 1

highQ - (빔)

첫 번째 수의 경우, lowQ에 일단 넣습니다. 또한, 지금까지 입력받은 숫자가 홀수개이므로 lowQ의 맨 위에 있는 1을 출력합니다.

 

2회차) 입력받은 수 2, 남은 수 - 6 9 2

lowQ의 맨 윗 값인 1 보다 2가 크므로 highQ에 넣습니다.

lowQ - 1

highQ - 2

 

3회차) 입력받은 수 6, 남은 수 - 9 2

입력받은 수 6이 lowQ의 맨 윗값인 1보다 크므로 highQ에 넣습니다.

lowQ - 1, highQ - 2, 6

하지만, 입력받은 수가 홀수개일 때, lowQ가 1만큼 더 크도록 해주도록 해야 하므로,

highQ에서 맨 위에 있는 수를 꺼내 lowQ에 넣습니다.

lowQ - 1, 2 highQ - 6

또, lowQ의 맨 위에 있는 값 (2)를 꺼냅니다. lowQ는 내림차순 입니다.

 

 

4회차) 입력받은 수 9, 남은 수 - 2

lowQ의 맨 위에 있는 수 2 보다 입력받은 수 9가 크므로 highQ에 쏙 넣어줍시다.

lowQ - 1, 2   highQ - 6, 9

 

5회차) 입력받은 수 11, 남은수 - x

lowQ의 맨 위에 있는 수 2 보다 입력받은 수 11이 크므로 highQ에 넣습니다.

lowQ - 1, 2, highQ - 6, 9, 11

다만, 이번에서 lowQ가 highQ 보다 1 만큼 그게 해 줄 것이므로, highQ 하나를 꺼내 lowQ에 하나 넣습니다.

lowQ - 1, 2, 6 highQ - 9, 11

또, 홀수번째 숫자이므로, lowQ 맨 위에 있는 숫자를 출력합니다.

 

위와 같은 방식으로, 매 번 중앙값을 찾을 수 있습니다.

 

 

 

소스코드

import java.lang.StringBuilder
import java.util.Collections
import java.util.PriorityQueue

fun main() = with(System.`in`.bufferedReader()) {
    val sb = StringBuilder()

    repeat(readLine().toInt()) {
        val n = readLine().toInt()

        sb.append((n + 1) / 2).append("\n")

        val lowQ = PriorityQueue<Int>(Collections.reverseOrder())
        val highQ = PriorityQueue<Int>()
        var cnt = 1
		
        // 10개씩 띄워서 주므로
        repeat(n / 10 + 1) { row ->
            val nums = readLine().split(" ").map { it.toInt() }

            for (col in nums.indices) {
            	// 10개 마다 한 줄 띄워서 출력
                // cnt는 매번 +1 해주고, 중앙값은 홀수일 때만 출력하므로 cnt%20 == 0
                if (cnt % 20 == 0) {
                    sb.append("\n")
                }

				// lowQ의 맨 윗값 보다 크다면 highQ에 offer
                if (lowQ.isNotEmpty() && nums[col] > lowQ.peek()) {
                    highQ.offer(nums[col])
                } else {
                // lowQ의 맨 윗값 보다 작다면 lowQ에 offer
                    lowQ.offer(nums[col])
                }

				// 두 개 큐 개수 맞춰주기
                while (lowQ.size < highQ.size) {
                    lowQ.offer(highQ.poll())
                }
                while (lowQ.size - 1 > highQ.size) {
                    highQ.offer(lowQ.poll())
                }

				// 홀수개일 경우 출력
                if (cnt++ % 2 == 1) {
                    sb.append(lowQ.peek()).append(" ")
                }
            }
        }

        sb.append("\n")
    }
    print(sb)
}

 

 

 

 

후기

비슷한 문제를 한 번 풀어봤기에, 그리 어렵지 않게 풀 수 있었습니다.

우선순위 큐를 이렇게 쓴다고? 하면서 풀었던 기억이 나네요 ㅎㅎ...