[백준 3163번] 떨어지는 개미 (Python3)

import sys,collections
input = sys.stdin.readline

for _ in range(int(input())):
  N,L,k = map(int,input().split())
  location,id = [],collections.deque()
  fall = []
  for _ in range(N):
    l,i = map(int,input().split())
    location.append(l)
    id.append(i)
    if i < 0:
      fall.append((location.pop(),id.popleft()))
  while location:
    fall.append((L-location.pop(),id.pop()))
  print(sorted(fall)[k-1][1])

스택을 이용해서 해결하였다.

아이디어를 떠올리는게 굉장히 중요한 문제이다.

 

일단 풀이를 요약하면,

1. 개미의 좌표와 ID를 따로 저장한다. (이때 ID는 popleft를 해야하므로 deque에 저장하는것이 좋다)

2. 만약 ID가 양수라면 (즉 오른쪽으로 이동한다면) 그대로 저장하고 넘어간다.

3. 만약 ID가 음수라면 (즉 왼쪽으로 이동한다면) location의 마지막 값과 id의 첫번째 값을 pop한 후 fall에 저장한다. (이 값들은 왼쪽으로 떨어지는 개미들이다)

4. 모두 끝나면, 남아있는 개미들을 pop한 뒤, fall에 저장한다. (오른쪽으로 떨어지는 개미들이다)

5. fall을 sort한 뒤 k번째 값을 출력한다.

 

이제 이렇게 되는 이유를 설명해야 하는데, 원리는 간단한데 설명하기가 쉽지 않다.

일단 설명해보자면, 

그림을 보면 현재 5,8에 위치한 개미 +4,+5가 오른쪽으로 이동하고 있다.

이때 왼쪽으로 이동하는 -1 개미가 등장하면, +5 개미와 -1 개미는 충동할 것이다. 그럼 -1 개미는 오른쪽으로, +5 개미는 왼쪽으로 이동하게 된다. 그리고 +5 개미는 +4개미와 충돌할 것이고, +4 개미는 왼쪽으로, +5 개미는 오른쪽으로 이동할 것이다. 그리고 +4 개미의 왼쪽에는 아무것도 없으므로 +4 개미는 왼쪽으로 떨어진다.

즉 최종적으로, 오른쪽으로 이동하는 개미들이 있을 때, 왼쪽으로 이동하는 개미가 한마리가 들어오면, 가장 왼쪽의 개미가 한마리 떨어지고, 나머지 개미들은 오른쪽으로 이동한다. 그래서 3번에서 id가 음수인 개미가 들어오면 id의 첫번째 값을 pop하는 것이다.

그렇다면 왜 location은 마지막 값을 pop 하는지의 이유도 알아야 하는데, 이는 직접 계산해보는게 아마 빠를 것이다. 충돌을 전부 계산해 보면, 가장 왼쪽에 있는 +4 개미는 19초 (즉 -1 개미의 위치) 가 지나면 떨어지게 된다. 충돌이 일어나면서 한칸씩 밀리는 것이다. 

그리고 모든 충돌이 끝나면 나머지 개미들은 전부 오른쪽으로 떨어지게 된다.

 

아이디어가 재미있고, 구현하는것도 재미있는 좋은 문제였다.

 

 

오늘의 교훈) 다양한 애드 혹 문제를 풀어보자