[Kotlin] 백준 28078 : 중력 큐

문제 링크 : https://www.acmicpc.net/problem/28078

문제 해설

조금 변형된 큐 문제. 큐의 방향에 따라 내용을 처리해야 한다.

 

큐의 방향에 따라서 공이 앞부분으로 나올 수도 있고 뒤부분으로 나올 수도 있기 때문에 덱(Deque)을 사용한다. 편의 상 본문에서는 큐라고 언급한다.

 

큐의 방향에 대해 생각해보자. 큐의 방향은 90도씩 회전하는 4가지 방향만 존재하기 때문에 큐의 앞부분이 바라보는 방향을 오른쪽부터 시계방향으로 0, 1, 2, 3으로 지정할 수 있다.

var direction = 0 // 0: right, 1: down, 2: left, 3: up

각 쿼리에 대한 동작을 구현하기 전에 큐에서 공의 개수와 가림막의 개수를 구하는 방법을 생각해보자.

 

단순히 count() 메소드를 사용해서 공 또는 가림막의 개수를 구하게 된다면 개수를 구할 때마다 큐 전체를 순회해야 한다. 쿼리의 개수가 최대 50만개인데, 이 방법을 사용했을 때 25만개의 값을 push하고 25만번의 count를 하게 되면 \(250,000\times 250,000=\) 625억번의 연산이 필요하기 때문에 시간초과가 발생한다.

 

시간초과를 방지하기 위해 공의 개수를 저장할 ballCount와 가림막의 개수를 저장할 wallCount를 변수로 선언해서 공이나 가림막이 추가/제거될 때마다 개수를 증가/감소 시키고 count 쿼리가 실행되면 해당 값을 출력하기만 하면 된다.

 

이제 각 쿼리의 동작을 구현해보자.

---

push

문제의 조건을 보면 큐가 세로로 세워져 있을 때는 공만 쏟아진다. 반대로 생각해보면 가림막은 큐의 방향에 상관 없이 push된다는 것을 알 수 있다.

if (char == 'w') {
    queue.addLast(char)
    wallCount++
    return@repeat
}

공을 집어넣는 경우는 큐의 방향에 따라 다르게 구현해야 한다.

 

• 큐의 앞방향이 아래를 향하고 있을 때(= 뒷방향이 위를 향하고 있을 때)
  이 때는 큐가 세로로 세워져 있기 때문에 큐의 가장 앞부분이 가림막으로 막혀있을 때만 공을 넣을 수 있다. 큐의 가장 앞부분에 가림막이 없다면 공을 넣자마자 빠지기 때문에 별도로 동작을 수행할 필요가 없다.
• 큐의 앞방향이 위를 향하고 있을 때(= 뒷방향이 아래를 향하고 있을 때)
  이 경우는 아래 그림처럼 가림막의 유무에 관계 없이 공을 넣자마자 빠지기 때문에 아무 동작을 수행하지 않아도 된다.
  

  

  큐의 앞방향이 위로 향한다면 공을 넣자마자 쏟아지게 된다

• 큐가 가로로 놓여있을 때
  이 때는 공이 쏟아지지 않으므로 평범하게 공을 추가하기만 하면 된다.

val char = nextToken()[0]
if (char == 'w') {
    queue.addLast(char)
    wallCount++
    return@repeat
}
when (direction) {
    1 -> if (queue.firstOrNull() == 'w') {
        queue.addLast(char)
        ballCount++
    }
    0, 2 -> {
        queue.addLast(char)
        ballCount++
    }
}

---

pop

큐의 가장 앞에 있는것을 빼내고 빼낸 물건의 종류에 따라 개수를 감소시키면 된다.

 

여기서도 큐가 세로로 세워져있는 경우를 고려해야 한다.

 

• 큐의 앞방향이 아래를 향하고 있을 때
  큐의 앞방향이 아래를 향하고 있다면 아래 그림처럼 가림막 밑에 있는 공이 모두 쏟아지게 된다. 즉, 큐의 가장 앞 요소가 공이 아닐 때까지 큐의 요소를 계속 빼내야 한다.
  

  

• 큐의 앞방향이 위를 향하고 있을 때
  큐의 앞방향이 위를 향하고 있다면 큐 가장 아래쪽의 가림막은 아무런 영향을 받지 않기 때문에 추가 동작을 수행할 필요가 없다.
  

  

if (queue.isEmpty()) return@repeat
if (queue.first() == 'w') wallCount-- else ballCount--
queue.removeFirst()
if (direction == 1) {
    while (queue.firstOrNull() == 'b') {
        queue.removeFirst()
        ballCount--
    }
}

---

rotate

큐의 방향을 반시계방향 또는 시계방향으로 회전시킨다.

 

반시계방향으로 회전하면 direction에 3을 더한 후, 시계방향으로 회전하면 direction에 1을 더한 후 direction을 4로 나눈 나머지값을 취해준다. 반시계방향 회전 연산은 음수 연산을 방지하기 위해 3을 더하는 방식으로 구현한다.

 

큐를 회전했을 때 세로로 세워지면 가림막 아래의 공이 모두 쏟아져야 한다.

• 큐의 앞방향이 아래를 향하고 있을 때(= 큐의 뒷방향이 위를 향하고 있을 때)
  큐의 가장 앞에 공이 나오지 않을 때까지 큐의 모든 요소를 앞에서 빼내고 공의 개수를 감소시킨다.
• 큐의 앞방향이 위를 향하고 있을 때(= 큐의 뒷방향이 아래를 향하고 있을 때)
  큐의 가장 뒤에 공이 나오지 않을 때까지 큐의 모든 요소를 뒤에서 빼내고 공의 개수를 감소시킨다.

direction = if (nextToken() == "l") (direction + 3) % 4 else (direction + 1) % 4
if (queue.isEmpty()) return@repeat
when (direction) {
    1 -> while (queue.firstOrNull() == 'b') {
        queue.removeFirst()
        ballCount--
    }
    3 -> while (queue.lastOrNull() == 'b') {
        queue.removeLast()
        ballCount--
    }
}

---

count

쿼리 종류에 따라 각각 ballCount 또는 wallCount를 출력하면 된다.

sb.append(if (nextToken() == "b") ballCount else wallCount).append('\n')

Code

import java.util.StringTokenizer

fun main() = with(System.`in`.bufferedReader()) {
    val sb = StringBuilder()
    val queue = ArrayDeque<Char>()
    var direction = 0 // 0: right, 1: down, 2: left, 3: up
    var ballCount = 0
    var wallCount = 0
    repeat(readLine().toInt()) {
        StringTokenizer(readLine()).apply {
            when (nextToken()) {
                "push" -> {
                    val char = nextToken()[0]
                    if (char == 'w') {
                        queue.addLast(char)
                        wallCount++
                        return@repeat
                    }
                    when (direction) {
                        1 -> if (queue.firstOrNull() == 'w') {
                            queue.addLast(char)
                            ballCount++
                        }
                        0, 2 -> {
                            queue.addLast(char)
                            ballCount++
                        }
                    }
                }
                "pop" -> {
                    if (queue.isEmpty()) return@repeat
                    if (queue.first() == 'w') wallCount-- else ballCount--
                    queue.removeFirst()
                    if (direction == 1) {
                        while (queue.firstOrNull() == 'b') {
                            queue.removeFirst()
                            ballCount--
                        }
                    }
                }
                "rotate" -> {
                    direction = if (nextToken() == "l") (direction + 3) % 4 else (direction + 1) % 4
                    if (queue.isEmpty()) return@repeat
                    when (direction) {
                        1 -> while (queue.firstOrNull() == 'b') {
                            queue.removeFirst()
                            ballCount--
                        }
                        3 -> while (queue.lastOrNull() == 'b') {
                            queue.removeLast()
                            ballCount--
                        }
                    }
                }
                "count" -> sb.append(if (nextToken() == "b") ballCount else wallCount).append('\n')
            }
        }
    }
    println(sb)
}