[코딩테스트] 백준 수열(2559) 파이썬(Python)

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

2559번: 수열

이 문제를 보자마자 완전탐색으로 풀어야겠다는 마음으로 다음의 코드로 제출했더니 시간초과로 통과하지 못하였습니다.

n, k = map(int, input().split())
array = list(map(int, input().split()))
result = 0
for s in range(n-k+1):
    e = s+k
    if sum(array[s:e]) > result:
        result = sum(array[s:e])
print(result)

저는 시간초과를 해결하기 위하여 누적합을 구하여 문제를 해결하였습니다.

n, k = map(int, input().split())
array = list(map(int, input().split()))
prefix_sum = [0]
result = []
sum_value = 0
for i in range(n):
    sum_value += array[i]
    prefix_sum.append(sum_value)

for i in range(n-k+1):
    result.append(prefix_sum[i+k] - prefix_sum[i])
    
print(max(result))

n=10, k=5

주어진 배열(array) : 3 -2 -4 -9 0 3 7 13 8 -3

이라고 할 때 부분 수열의 합을 구해보면서 문제해결과정을 설명하겠습니다.

 

array의 인덱스와 값을 표현해보면 다음과 같습니다.

0	1	2	3	4	5	6	7	8	9
3	-2	-4	-9	0	3	7	13	8	-3

인덱스를 기준으로 0부터 4까지의 합을 구하고 싶다면 sum(array[0:5])하면 구할 수 있습니다.

1 ~ 5 : sum(array[1:6])

2 ~ 6 : sum(array[2:7])

...

i ~ i+k : sum(array[i:i+k-1])

이렇게 구하면 완전탐색으로 구하는 것이기 때문에 시간초과에 걸리게 됩니다.

 

식을 약간 바꿔보겠습니다.

1 ~ 5 : sum(array[:6]) - sum(array[:1])

2 ~ 6: sum(array[:7]) - sum(array[:2])

3 ~ 7 : sum(array[:8]) - sum(array[:3])

...

i ~ i+k : sum(array[:i+k]) - sum(array[:i])

 

2~6까지의 합을 구하고 싶다면 7까지 더한 값에서 2까지 더한 값을 빼면 구할 수 있습니다. 쉽게 표현하자면 다음과 같습니다.

[3 -2 -4 -9 0 3 7 13 8 -3]

주황색 부분에서 분홍색 부분을 제거하면 구할 수 있습니다.

3 ~ 7까지의 합이라면

[3 -2 -4 -9 0 3 7 13 8 -3]

이 됩니다.

 

하지만 매번 저 부분 수열을 구하고 모두 더하여 값을 구하면 역시 시간초과가 나오게 됩니다.

그래서 부분 수열을 모두 더한 값을 보관하고 있는 prefix_sum 이 필요한 것입니다.

 

prefix_sum의 가장 첫번째 원소[0]는 [0] 입니다.

prefix_sum[1] : sum(array[:1])(1까지의 합)

prefix_sum[2] : sum(array[:2])(2까지의 합)

prefix_sum[n-1] : sum(array[:n-1])(n-1까지의 합)

 

하지만 코드를 보면 sum(array[:x])로 값을 구하지 않고 sum_value의 값을 저장하여 구하고 있습니다.

sum(array[:x])로 부분수열을 찾아 더하다보면 역시 시간초과가 나오게 됩니다.

 

prefix_sum으로 누적합을 구하여 매번 배열의 합을 구하는 과정을 최소화하여 시간초과를 통과하고 문제를 해결할 수 있었습니다.