[Baekjoon] 14002/구간 합 구하기 4/Python/파이썬/누적 합

💡문제

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

입력

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

출력

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

예제입력

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

예제출력

12
9
1

---

📖내가 작성한 Code

import sys

''''
구간 합 연습 문제
'''


def make_prefix_sum(N, lst):
    prefix_sum = [0] * (N + 1)
    for i in range(1, N+1):
        prefix_sum[i] = prefix_sum[i - 1] + lst[i -1]
    return prefix_sum


def find_prefix_sum(prefix_sum, panels):
    return (prefix_sum[panel[1]] - prefix_sum[panel[0] - 1] for panel in panels)


def main():
    (N, M), lst, *panels = [list(map(int, line.split())) for line in sys.stdin.read().splitlines()]
    print(*find_prefix_sum(make_prefix_sum(N, lst), panels), sep='\n')


if __name__ == '__main__':
    main()

---

✍️풀이과정

누적 합 연습을 위해서 스무스하게 푼 문제.

그런데, 어느 분께서 푸신 코드를 보고 놀랬다.

from itertools import*
from sys import*
if __name__ == "__main__":
    Q=map(str.split, stdin.read().splitlines())
    next(Q)
    A = list(accumulate(map(int, next(Q)),initial=0))
    stdout.write('\n'.join(f'{A[int(b)]-A[int(a)-1]}'for a,b in Q))

next를 너무 깔끔히 쓰시는 것을 보고 나도 다음에는 도전해봐야겠다는 생각이 들었다. 코드보고 아름답다라는 생각 오랫만이었음.

---

🧠 코드 리뷰

1. 파이썬 내장 모듈:
   파이썬 내장 모듈인 itertools.accumulate를 활용하면 코드가 더욱 간결해질 수 있습니다.

---

🛠️AI 개선 코드

import sys
from itertools import accumulate

def main():
    input_data = sys.stdin.read().splitlines()
    N, M = map(int, input_data[0].split())
    numbers = list(map(int, input_data[1].split()))
    prefix_sum = [0] + list(accumulate(numbers))
    output = []
    for line in input_data[2:]:
        i, j = map(int, line.split())
        output.append(str(prefix_sum[j] - prefix_sum[i - 1]))
    sys.stdout.write("\n".join(output))

if __name__ == '__main__':
    main()

---

💻결과

백준문제 보러가기

---

🖱️참고 링크

누적합 알고리즘