새소식

PS/백준

[백준/3163] 떨어지는 개미 (Python)

  • -

Problem : https://www.acmicpc.net/problem/3163

 

3163번: 떨어지는 개미

첫째 줄에 테스트 케이스의 개수 T가 주어진다. 각 테스트 케이스의 첫째 줄에는 N, L, k가 주어진다. 다음 N개 줄에는 pi와 ai가 주어진다. ai는 개미의 ID이고, pi는 그 개미의 초기 위치이다. 항상 p

www.acmicpc.net

 

Difficulty : Platinum 5

 

Status : Solved

 

Time : 01:44:32

 


 

문제 설명

 

더보기

 

개미 N마리가 막대 위에 올라가 있다. 일부 개미는 왼쪽을 바라보고 있고, 나머지 개미는 오른쪽을 바라보고 있다. 모든 개미는 매우 작아서 크기가 없는 점으로 나타낼 수 있다. 시작 신호가 주어지면, 개미는 바라보고 있는 방향으로 행진을 시작한다. 모든 개미는 동일한 속도 초속 1mm로 이동한다. 두 개미가 한 점에서 충돌하는 경우가 발생할 수 있다. 이 경우에 두 개미는 행진하는 방향을 반대 방향으로 바꾸고, 행진을 계속하게 된다. 개미가 방향을 바꾸는데 걸리는 시간은 없다. 개미가 막대의 끝에 도착하는 경우에는, 막대에서 떨어지게 된다. 막대는 땅 위에 떠있다고 가정한다.

처음에 모든 개미의 위치는 서로 다르다. 즉, 두 개미가 막대 위의 한 점에 같이 있는 경우는 없다. 개미는 부호 있는 정수로 나타낼 수 있다. 이 정수를 개미의 ID라고 한다. 개미의 ID의 부호는 개미가 처음에 바라보고 있는 방향이다. -는 왼쪽을 바라보고 있는 것이고, +는 오른쪽을 바라보고 있는 것이다. 개미의 ID의 절댓값은 1부터 109까지의 정수 중 하나이다. 또, 모든 개미의 ID의 절댓값은 서로 다르다. 아래 그림에는 개미가 총 6마리가 있고, ID는 {+4, +5, -1, -3, -2, +6}이다. 각 개미의 초기 위치는 {5, 8, 19, 22, 24, 25}이며, 막대의 길이 L = 30이다. 화살표는 처음에 개미가 바라보고 있는 방향을 나타낸다. 왼쪽 끝의 좌표는 0이고, 오른쪽 끝의 좌표는 30이다. ID가 +6인 개미는 시간 t = 5일 때, 막대의 오른쪽 끝에 도착하며, t = 6에 막대에서 떨어지게 된다.

 

 

개미가 행진을 시작하기 전의 상태 (ID와 막대 상의 위치)가 주어진다. 두 개미가 동시에 막대의 양 끝에서 떨어지는 경우에는, ID가 작은 개미가 조금 더 먼저 떨어진다고 한다. 아래 그림은 이와 같은 경우를 나타낸 그림이다. 두 개미 {-1, +2}는 끝에 동시에 도착하게 된다. -1 < +2 이기 때문에, ID가 -1인 개미가 +2인 개미보다 조금 더 먼저 떨어지게 된다. 따라서, 아래 그림의 네 개미가 떨어지는 순서는 {-1, 2, 4, 3}이 된다.

 

양의 정수 1 ≤ k ≤ n이 주어졌을 때, k번째로 떨어지는 개미를 구하는 프로그램을 작성하시오.

 

 

입력 및 출력

 

더보기

입력

 

첫째 줄에 테스트 케이스의 개수 T가 주어진다. 각 테스트 케이스의 첫째 줄에는 N, L, k가 주어진다. 다음 N개 줄에는 pi와 ai가 주어진다. ai는 개미의 ID이고, pi는 그 개미의 초기 위치이다. 항상 pi가 증가하는 순서로 (pi<pi+1) 주어진다. (1 ≤ pi ≤ L-1, 3 ≤ N ≤ 100,000, 10 ≤ L ≤ 5,000,000, 1 ≤ k ≤ N)

 

출력

 

각 테스트 케이스마다, N마리 개미 중에서 k번째로 떨어지는 개미의 ID를 출력한다. 개미의 ID가 양수인 경우에 +를 출력하면 안 된다.

 

입력 예시

 

2
6 30 3
5 4
8 5
19 -1
22 -3
24 -2
25 6
4 35 2
5 -1
12 3
20 4
30 2

 

출력 예시

 

-2
2

 

 


 

풀이

 

꽤 많이 고민을 했던 문제. 우선 어떻게 하든 개미의 인덱스 순서는 변하지 않는다는 점에 착안해서, 각 개미가 충돌할 때마다 최종 충돌 시간, 방향과 위치를 업데이트해 주었다. 즉 모든 개미의 방향이 '<<<<<>>>>>>'꼴이 된다면, 충돌 위치, 그리고 최종 충돌 시간을 토대로 개미들의 추락 순서를 업데이트해 줄 수 있다고 보았다.

 

import sys
input = sys.stdin.readline

def cmp(x) :
  if x[1] < 0:
    return (x[0] + x[2], x[3])
  else :
    return (L - x[0] + x[2], x[3])
    
def solve() :
  global L
  N, L, K = map(int, input().split())
  ants = list()
  for _ in range(N) :
    p, a = map(int, input().split())
    ants.append([p, a, 0, a])
  while True :
    conflict = False
    for i in range(N-1) :
      if ants[i][1] > 0 and ants[i+1][1] < 0 :
        conflict = True
        max_t = max(ants[i][2], ants[i+1][2])
        tvar = abs(ants[i][2]-ants[i+1][2])
        moved = (ants[i+1][0] - ants[i][0] - tvar) / 2
        if max_t == ants[i][2] :
          ants[i][0] = ants[i+1][0] = ants[i][0] + moved
        else :
          ants[i][0] = ants[i+1][0] = ants[i+1][0] - moved
        ants[i][2] = ants[i+1][2] = max_t + moved
        ants[i][1] = -ants[i][1]
        ants[i+1][1] = - ants[i+1][1]
    if not conflict :  
      break
  ants.sort(key = cmp)
  print(ants[K-1][-1])

for _ in range(int(input())) :
  solve()

 

...문제는 개미의 충돌 시점을 계속해서 계산하고 업데이트하므로, 최악의 경우 비선형적인 시간복잡도가 소요될 수 있다는 점이었다. 즉 이 경우는 무조건 시간 초과를 일으킨다.

 

두 번째 관찰 결과, 개미들의 이동량은 충돌 여부와 상관없이 보존된다는 점을 발견했다. 즉 개미들이 충돌할 때, 실제로는 "인덱스 정보만을 서로 넘겨주면서 그대로 이동"하는 것과 다를바 없다. 이를테면

 

 

꼴로 이동한다고 했을때, 먼저 (+5, -1)이 충돌하고 방향은 > < > 꼴이 된다. 그 다음으로는 (+4, +5) 가 충돌하고 방향은 <>>꼴이 될 것이다. 즉 -1에서 19만큼 이동하는 이동량은 최종적으로 +4에 전달되며, +4의 이동량은 +5, +5의 이동량은 -1로 전달된다. 이동량이 회전하는 셈이다. 또한 이동 방향 역시 끝단으로 전달됨을 알 수 있다. 즉 단 한번의 순회로 모든 이동량 전달에 대해 업데이트가 가능해진다.

 

즉 순회 중 현재 개미가

  • 순회 방향과 동일하다면 그대로 큐에 저장하도록 한다.
  • 그렇지 않다면, 큐의 왼쪽(제일 오래된 idx) 에 존재하는 개미를 탈출시키고, 그 탈출하는 개미의 이동량은 현재 개미의 이동량과 같다. 즉 큐에 저장한 후, idx는 popleft, 이동량은 pop을 진행하면 된다.

이런 식으로, O(N)의 시간복잡도 내에 모든 개미와 최종 이동거리를 배치할 수 있게 된다.

 

풀이 코드

from collections import deque
import sys
input = sys.stdin.readline

def solve() :
  global L
  N, L, K = map(int, input().split())
  idx_q = deque()
  fall_q = deque()
  result = list()
  for _ in range(N) :
    p, idx = map(int, input().split())
    if idx > 0 :
      p = L - p
    idx_q.append(idx)
    fall_q.append(p)
    if idx < 0  :
      result.append((fall_q.pop(), idx_q.popleft()))
  while idx_q and fall_q :
    result.append((fall_q.pop(), idx_q.pop()))
  result.sort()
  print(result[K-1][-1])

for _ in range(int(input())) :
  solve()

 

풀이 완료!

Contents

포스팅 주소를 복사했습니다

이 글이 도움이 되었다면 공감 부탁드립니다.