BOJ 11441번 : 합 구하기 (Python/Silver 3)

11441번: 합 구하기

문제

N개의 수 A1, A2, ..., AN이 입력으로 주어진다. 총 M개의 구간 i, j가 주어졌을 때, i번째 수부터 j번째 수까지 합을 구하는 프로그램을 작성하시오.

 

입력

첫째 줄에 수의 개수 N이 주어진다. (1 ≤ N ≤ 100,000) 둘째 줄에는 A1, A2, ..., AN이 주어진다. (-1,000 ≤ Ai ≤ 1,000) 셋째 줄에는 구간의 개수 M이 주어진다. (1 ≤ M ≤ 100,000) 넷째 줄부터 M개의 줄에는 각 구간을 나타내는 i와 j가 주어진다. (1 ≤ i ≤ j ≤ N)

 

출력

총 M개의 줄에 걸쳐 입력으로 주어진 구간의 합을 출력한다.

---

예제 입력

5
10 20 30 40 50
5
1 3
2 4
3 5
1 5
4 4

예제 출력

60
90
120
150
40

---

문제 설명

N개의 수를 i ~ j까지 합한 결과를 구한다.

 

풀이

- 누적 합 알고리즘 이용

 

누적 합 알고리즘

array = [10, 20, 30, 40, 50]
n = len(array)
prefix_sum = [0] * (n + 1)

for i in range(n):
    prefix_sum[i + 1] = prefix_sum[i] + array[i]
   
#결과 : prefix_sum[0, 10, 30, 60, 100, 150]

누적 합 알고리즘 설명 - https://dev-cloud.tistory.com/141

 

파이썬 출력 속도 개선 방법

※  sys.stdout.write( )

• 줄바꿈 없이 바로 이어서 출력
• 메모리 큼 / 속도 빠름
• () 안에 문자열이 들어가야 하므로 정수일 경우는 str(정수) 해줘야 한다.
• 줄바꿈 하면서 연속 출력 하고 싶다면 끝에 +'\n'을 넣어주자.

Solution

import sys
input = sys.stdin.readline
N = int(input())
A = list(map(int, input().split()))

prefix_sum = [0] * (N + 1)
for k in range(N):
    prefix_sum[k+1] = prefix_sum[k] + A[k]
    
M = int(input())
ans = 0
for _ in range(M):
    i, j = map(int, input().split())
    ans = prefix_sum[j]-prefix_sum[i-1]
    sys.stdout.write(str(ans) + '\n')

시간 초과 풀이

#1. 시간 초과 생각 안 하고 기본적으로 생각할 수 있는 풀이지만 당연히 시간 초과
for _ in range(M):
    i, j = map(int, input().split())
    result = 0
    while i <= j:
        result += A[i-1]
        i += 1
    print(str(result))
    
#2. 이중 for문은 당연히 O(n^2) 시간이 걸리니 시간초과가 날 수 밖에..
for _ in range(M):
    i, j = map(int, input().split())
    prefix_sum = [0] * (j+1)
    for k in range(i, j+1):
        prefix_sum[k] = prefix_sum[k-1] + A[k-1]
    print(prefix_sum[j])
    
#3. 누적 합 알고리즘을 썼는데도 시간 초과가 나길래 억울할 지경
prefix_sum = [0] * (N + 1)
for k in range(N):
    prefix_sum[k+1] = prefix_sum[k] + A[k]

for _ in range(M):
    i, j = map(int, input().split())
    print(prefix_sum[j]-prefix_sum[i-1])​