[백준] 구간 합 구하기 4 - 11659번 (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트 제출

import sys
input_sys = sys.stdin.readline

# 11659번
n, m = map(int, input_sys().split())
if(not(1 <= n <= 100000) or not(1 <= m <= 100000)):
    exit()
num = list(map(int, input_sys().split()))

for _ in range(m):
    ans = 0
    i, j = map(int, input_sys().split())
    if(not(1 <= i <= j <= n)):
        exit()
    for k in range(i-1, j):
        ans += num[k]
    print(ans)

수정 제출

import sys
input_sys = sys.stdin.readline

# 11659번
n, m = map(int, input_sys().split())
if(not(1 <= n <= 100000) or not(1 <= m <= 100000)):
    exit()
num = list(map(int, input_sys().split()))

# 누적합 구하기
sums = [0] * (n + 1)
for i in range(1, n + 1):
    sums[i] = sums[i-1] + num[i-1]

for _ in range(m):
    i, j = map(int, input_sys().split())
    if(not(1 <= i <= j <= n)):
        exit()
    print(sums[j] - sums[i-1])

풀이

1트 제출 설명

 

python의 경우 input() 함수는 입력을 받을 때마다 버퍼를 비우고 입력을 처리하기 때문에 느리게 동작하게 됩니다. 그렇기 때문에 sys.stdin,readline()을 사용하여 빠르게 입력받는 형식을 채택하게 되었고, 이를 항상 기본적인 조건으로 사용하게 됩니다.

import sys
input_sys = sys.stdin.readline

n과 m에 각각의 정수값을 입력받아 대입하게 됩니다. n은 수의 개수와 m은 합을 구해야 하는 구간입니다. 이후 제한에 걸려있는 것을 해주기 위하여 n 혹은 m이 1보다 크거나 같지 않고 100,000보다 작거나 같지 않으면 코드가 바로 종료되게 하였습니다.

num 이란 변수에 값을 넣을건데 입력된 값을 리스트 형식으로 만들어 num에 저장하게 됩니다.

이후 m번 만큼 반복을 하며 값 2개를 입력받아 각각 i와 j에 대입을 해줍니다. 만약 i와 j가 아래 if문에 부합하지 않는다면 코드는 종료가 됩니다.

조건문을 넘기게 된다면 i-1 값부터 j 값까지 k라는 변수에 값이 넣어지며 반복이 됩니다. 반복이 되는 동안 ans라는 변수에는 num[i-1] 값부터 num[j-1] 값까지 전부 더해지게 되고 끝난다면 ans에 저장된 값이 출력이 되게 됩니다.

이 코드는 출력값은 정답이 나오지만 시간 초과. 즉, 시간 복잡도에 의해 오답 처리가 되었고 시간 복잡도를 생각하며 코드를 다시 작성해보았습니다.

 

수정 코드 설명

4~8행 : 이 부분은 수정하기 전과 코드가 다르지 않아 간단히만 설명하고 넘어가겠습니다.

값 2개를 띄워쓰기로 나눠 입력받아 n과 m에 각각 대입하여 줍니다. 이후 조건문을 이용하여 n과 m이 조건에 부합하는지 검토하고 부합하지 않다면 코드는 종료됩니다.

다음 행에는 값을 n개 입력을 받게 되는데 이것을 리스트 형식으로 변수 num에 저장하게 됩니다.

# 11659번
n, m = map(int, input_sys().split())
if(not(1 <= n <= 100000) or not(1 <= m <= 100000)):
    exit()
num = list(map(int, input_sys().split()))

10~13행 :0이란 값이 저장되어 있는 배열을 총 n+1개 만들어냅니다. 이것을 sums에 저장합니다. 이후 1부터 n+1만큼의 값을 i에 대입하며 sums[i-1] + num[i-1]의 값을 sums[i]에 저장해줍니다. 이렇게 된다면 sums 에는 [0, 5, 9, 12, 14, 15] 가 저장되게 됩니다. 이것은 num에 입력하여 저장된 값을 순차적으로 하나씩 더한 값과 같습니다.

ex) num에 5 4 3 2 1 이 입력되었다면 5, 5+4, 5+4+3, 5+4+3+2, 5+4+3+2+1 이 sums에 저장된 것입니다.

# 누적합 구하기
sums = [0] * (n + 1)
for i in range(1, n + 1):
    sums[i] = sums[i-1] + num[i-1]

15~19행 : m번 만큼 반복이 됩니다. 반복이 될 때마다 값 2개를 띄워쓰기를 통해 분류하여 각각 i와 j에 대입해주며 조건문을 사용하여 조건에 부합하는지 확인합니다. 부합하지 않다면 코드는 종료되게 됩니다.

만약 조건에 부합하다면 sums[j] 에서 sums[i-1]를 빼낸 값을 출력하게 됩니다.

for _ in range(m):
    i, j = map(int, input_sys().split())
    if(not(1 <= i <= j <= n)):
        exit()
    print(sums[j] - sums[i-1])

 

수정된 코드는 배열의 누적합을 미리 구해둬 각 쿼리마다 sums 배열을 이용해 누적합을 계산하는 방법을 사용하였습니다. 이렇게 한다면 반복문이 최대 m번만 실행하게 되어 시간 복잡도를 줄일 수 있습니다.