[백준🥈3] #11659 구간 합 구하기 4 (Python)

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

11659번: 구간 합 구하기 4

 

문제

수 N개가 주어졌을 때, i번째 수부터 j번째 수까지 합을 구하는 프로그램을 작성하시오.

입력

첫째 줄에 수의 개수 N과 합을 구해야 하는 횟수 M이 주어진다. 둘째 줄에는 N개의 수가 주어진다. 수는 1,000보다 작거나 같은 자연수이다. 셋째 줄부터 M개의 줄에는 합을 구해야 하는 구간 i와 j가 주어진다.

출력

총 M개의 줄에 입력으로 주어진 i번째 수부터 j번째 수까지 합을 출력한다.

제한

• 1 ≤ N ≤ 100,000
• 1 ≤ M ≤ 100,000
• 1 ≤ i ≤ j ≤ N

예제 입력 1

5 3
5 4 3 2 1
1 3
2 4
5 5

예제 출력 1

12
9
1

 

❌ Code (시간초과)

import sys
input = sys.stdin.readline

n, m = map(int, input().split())
nums = [0] + list(map(int, input().split()))

for _ in range(m):
    i, j = map(int, input().split())
    print(sum(nums[i:j+1]))

✔️ Code

import sys
input = sys.stdin.readline

n, m = map(int, input().split())
nums = list(map(int, input().split()))
prefix_sum = [0] # 누적합을 저장할 배열 -> 인덱스를 편리하게 하기 위해 [0] 넣고 시작
sum = 0 # 누적합을 저장할 변수

for num in nums:
    sum += num  
    prefix_sum.append(sum) # 각 값들을 더해 배열에 넣어줌

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

✏️ Comment

드디어 풀어본 누적합 문제

시간초과로 실패한 코드는 O(NM)의 시간복잡도를 가지기 때문에 N과 M이 커지면 1초 안에 수행될 수 없다. 시간을 줄여주기 위해서 누적합 알고리즘을 사용했다.

prefix_sum 이라는 누적합을 저장할 배열을 하나 생성해서, nums에 있는 숫자들을 하나씩 더해주면서 prefix_sum에 append 해준다. 그리고 m개의 구간에 따라서 prefix_sum[j] - prefix_sum[i-1] 을 해주면 이 한 번의 연산으로 답을 쉽게 구할 수 있다.

prefix_sum[j] - prefix_sum[i-1] 을 수행하면 계산시간은 O(1)이 되고, 결과적으로 N개의 데이터와 M개의 입력이 있을 때, 전체 구간 합을 구하는 시간복잡도는 O(N+M)이 된다.