[Python] 백준 2560번: 짚신벌레

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

2560번: 짚신벌레

 

 > 접근

최초에는 [1] -> [0,1] -> [1,0,1] -> [1,1,0,1] -> [1,1,1,0,1] -> [2,1,1,1,0,1] -> [2,2,1,1,1,0] 식으로, 리스트를 right shift 시킨 뒤 a~b 범위의 합을 맨 앞에 추가하는 식으로 시도하였으나 시간 초과가 나왔다.

최적화 할 수 있는 부분은 총 두 가지가 있었다.

첫번째로 리스트를 right shift하고 맨 앞에 원소를 추가하는 작업을 최적화 시킬 수 있다. 먼저, 위에 제시한 리스트를 보면 알 수 있듯이, 뒷부분이 항상 같다. 따라서 이것을 뒤집어서 생각하기로 하였다. 이에 따라 리스트를 처음부터 N+1 크기로 정의하였으며, i와의 상대적인 위치로 a~b 구간을 정의한다.

두번째로 부분합을 구하는 과정이다. 무식하게 (a~b] 구간을 돌며 합을 구하는 것보단, prefix sum을 이용하여 앞뒤로 추가 및 제거를 해주는 것이 훨씬 효율적이다.

 

 > 풀이

다만 N이 작을 경우에는 예외가 나올 수 있다. 아래 예외 코드를 참고하자.

1 2 3 1

 

추가로 아래 코드를 빠르게 통과하면 TLE를 통과한다고 볼 수 있을 것이다.

1 9999 10000 1000000

 

# -*- coding: utf-8 -*-
import sys
input = sys.stdin.readline
DIV = 1000

if __name__ == '__main__':
  a, b, d, N = map(int, input().split())

  dp = [0]*(N+1)
  dp[0] = 1

  prefix_sum = 0
  for i in range(1, N+1):    
    prefix_sum = (prefix_sum + dp[i-a] - dp[i-b] + DIV) % DIV
    dp[i] = prefix_sum    
    
  sm = 0
  for i in range(max(0, N-d+1), N+1):
    sm = (sm + dp[i]) % DIV
  print(sm)