[Python] 백준 2943번 - 토끼

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

2943번: 토끼

풀이

일단 문제를 잘 이해해야한다. "그날 성냥을 넣은 성냥갑과 컵에 들어있는 성냥의 수를 구하는 프로그램"이라고 문제에 명시되어있지만, "그날 넣은 성냥"의 개수를 세는 것으로 착각하기 쉬운 문제이다. 그날 넣은 성냥의 개수를 세는 것으로 착각했다면 아마 출력 예시를 보고 당황했을 것이다. 

문제를 잘 읽어보고 헷갈릴 만한 내용들을 정리해보았다.

• m일간 컵과 성냥갑은 공유된다. 초기화 되는 것이 아니다.
• 같은 컵 뿐만 아니라 같은 성냥갑에도 성냥이 한 개 이상 들어갈 수 있다.
  (나는 성냥갑에는 성냥이 하나만 들어갈 것이라고 착각했었다.)
• 첫번째 날 부터 m번째 날 까지 각각의 날짜에 내가 성냥을 넣은 곳에 들어있는 성냥의 총 개수를 세는 것이다.

문제를 이해했으면 다음으로 어떻게 성냥갑에 넣을 성냥과 컵에 넣을 성냥을 구분할 것인지 생각해야한다. 일단 성냥갑은 토끼의 마리수 만큼 필요하고, 컵은 최대 [ 토끼수 / k + 1 ]개 만큼 필요하다. 단순히 성냥갑에 넣는것만 생각했을 때는 입력값 들어온 시작 순번부터 Index로 설정하고 구매한 딸기 수만큼 성냥갑[index] += 1을 반복해주면 된다.
이 과정 중에 성냥갑을 k개로 나눈 부분들 중, 한 부분의 시작지점 부터 끝 지점까지 전부 성냥을 넣는 경우가 발생하면 컵에 성냥을 넣어줘야한다. 

나는 딸기를 첫번째로 먹을 토끼의 start index와 마지막으로 먹을 토끼의 end index를 이용해서 세 부분으로 나누어 구현하였다. 위 그림에서 k개로 나눈 부분의 시작지점 == 어떤 컵의 시작구간 이라고 할 때 아래와 같이 세 부분으로 나누었다.

[ start ~ 어떤 컵의 시작구간 직전 ] - [ 어떤 컵의 시작구간 ~ 어떤 컵의 마지막구간 ] - [어떤 컵의 마지막 구간 직후 ~ end]

0부터 n - 1번까지의 토끼의 번호에 k*k <=  n을 만족하는 가장큰 k를 나머지 연산한 결과가 0이면 컵이 시작하는 지점이라고 할 수있다. 마찬가지로 토끼의 번호에 k를 나머지 한 연산 결과가 k - 1 이면 컵이 끝나는 지점이다. 이것을 이용해서 A, C 구간에서는 성냥갑에 성냥을, B구간에서는 컵에 성냥을 넣어주면 된다.

(아래의 코드에서 start는 0 부터 n - 1 범위, end는 1 부터 n의 범위로 두었기 때문에 컵이 끝나는 구간을 구할 때 end % k == k - 1이 아니라 end % k == 0 연산을 이용하는것에 주의하며 봐주세요.)

import sys
import math
input = sys.stdin.readline

n, m = map(int, input().split())

k = int(math.sqrt(n))
box = [0] * n
cup = [0] * (n // k + 1)

# 내가 한 타임에 "넣은 성냥 수"만 세는게 아니라 
# "성냥 넣은 곳들의 총 성냥"을 세는것임!!!
# 컵이든 성냥갑이든 성냥은 1개 이상 들어갈 수 있음!!

for _ in range(m):
    berry, start = map(int, input().split())

    # start는 0부터 N-1 순이지만 end는 1부터 N 순임.
    # 인덱싱이랑 컵 시작, 마지막 구간 계산할 때 주의 필요함
    start -= 1
    end = berry + start 
    ans = 0

    # 시작 토끼부터 컵의 시작 구간(index % k == 0) 전 까지 반복
    while (start < end and start % k != 0):
        box[start] += 1
        ans += box[start]
        start += 1

   # end != n 이라면 뒷 순번 토끼부터 확인하며 성냥갑에 성냥 넣어줌 
   # end % k == 0이 되면 컵의 마지막구간임. (start는 0부터 N-1 순이지만 end는 1부터 N 순임에 주의)
   # 컵의 마지막구간 만나기 전 까지 성냥갑에 성냥 넣어줌
    if (end != n):
        while (start < end and end % k != 0):
          end -= 1
          box[end] += 1
          ans += box[end]

  # 이제 start는 컵의 시작구간이고 end는 컵의 마지막구간임
  # 앞으로 넣을 성냥은 전부 컵에 들어감
    while(start < end):
        cup[start // k] += 1
        ans += cup[start // k]
        start += k

    print(ans)