[Baekjoon] 11660/구간 합 구하기 5/Python/파이썬/누적 합

井·2025년 2월 4일

PrefixSum python 누적합 백준 알고리즘 파이썬

문제풀이

목록 보기

28/56

💡문제

N×N개의 수가 N×N 크기의 표에 채워져 있다. (x1, y1)부터 (x2, y2)까지 합을 구하는 프로그램을 작성하시오. (x, y)는 x행 y열을 의미한다.

예를 들어, N = 4이고, 표가 아래와 같이 채워져 있는 경우를 살펴보자.

1 2 3 4
2 3 4 5
3 4 5 6
4 5 6 7

여기서 (2, 2)부터 (3, 4)까지 합을 구하면 3+4+5+4+5+6 = 27이고, (4, 4)부터 (4, 4)까지 합을 구하면 7이다.

표에 채워져 있는 수와 합을 구하는 연산이 주어졌을 때, 이를 처리하는 프로그램을 작성하시오.

입력

첫째 줄에 표의 크기 N과 합을 구해야 하는 횟수 M이 주어진다. (1 ≤ N ≤ 1024, 1 ≤ M ≤ 100,000) 둘째 줄부터 N개의 줄에는 표에 채워져 있는 수가 1행부터 차례대로 주어진다. 다음 M개의 줄에는 네 개의 정수 x1, y1, x2, y2 가 주어지며, (x1, y1)부터 (x2, y2)의 합을 구해 출력해야 한다. 표에 채워져 있는 수는 1,000보다 작거나 같은 자연수이다. (x1 ≤ x2, y1 ≤ y2)

출력

총 M줄에 걸쳐 (x1, y1)부터 (x2, y2)까지 합을 구해 출력한다.

예제입력

예제출력

27
6
64

📖내가 작성한 Code

import sys
import itertools

''''
2차원 배열합 구하기
'''


def main():
    query = map(str.split, sys.stdin.read().splitlines())
    N = int(next(query)[0])
    matrix = list(list(map(int, next(query)))for _ in range(N))
    prefix_sum = [[0]*(N+1)] + list(list(itertools.accumulate(matrix[i], initial=0)) for i in range(N))

    for i in range(1, N+1):
        for j in range(N+1):
            prefix_sum[i][j] += prefix_sum[i - 1][j]

    sys.stdout.write('\n'.join(
        str(prefix_sum[int(x2)][int(y2)] - prefix_sum[int(x1) - 1][int(y2)] - prefix_sum[int(x2)][int(y1) - 1] +
            prefix_sum[int(x1) - 1][int(y1) - 1])
        for x1, y1, x2, y2 in query
    ))


if __name__ == '__main__':
    main()

✍️풀이과정

이것 또한 단순한 누적합 문제. 어제 누적합 이쁜 코드를 보고 나도 좀 효율적으로 짜보자 생각해서 next도 써보고 print 대신 sys.stdout.write도 사용해봄.

여기서

prefix_sum = [[0]*(N+1)] + list(list(itertools.accumulate(matrix[i], initial=0)) for i in range(N))

에 문제가 조금 있었는데, list()와 '[]'의 차이였다.

[[0]*(N+1)]

리스트 리터럴(literal)의 표현식

리스트를 직접 정의

list([0]*(N+1))

list() 생성자를 사용해 [0]*(N+1)으로 생성된 리스트를 얕은 복사(shallow copy)한 결과를 반환한 표현식

다른 iterable을 리스트로 변환 시켜줌

따라서, [0]*(N+1)는 이미 리스트라서 list()로 감싸면 그냥 앝은 복사만 해서 결과를 반환해줌

그리고 sys.stdout.write를 사용해봤는데, print보다 빠른 이유는

print는 여러 가지 옵션 처리로 추가 오버헤드가 있다.

자동으로 개행 문자 추가, 여러 인자 처리와 문자열 변환

print는 end, sep, file, flush 등의 옵션을 처리하는 로직이 추가

근데 굳이 이렇게 처리해야하나 싶긴함.

그런데,

그런데 기대도 안했는데 맞힌 사람 python3 시간부분 in 3 해버림. 파이썬을 탐구하면서 이런 소소한 부분에서 희열을 느끼는 듯?

🧠 코드 리뷰

쿼리 입력 사전 처리: 입력 시 쿼리 내의 각 값들을 정수형으로 미리 변환해두면, 쿼리 처리 시 불필요한 반복적 형변환을 줄일 수 있습니다.

🛠️AI 개선 코드

import sys
import itertools

def main():
    input_lines = sys.stdin.read().splitlines()
    it = iter(input_lines)
    N, M = map(int, next(it).split())
    matrix = [list(map(int, next(it).split())) for _ in range(N)]
    
    prefix_sum = [[0] * (N + 1)]
    for i in range(N):
        prefix_sum.append(list(itertools.accumulate(matrix[i], initial=0)))
    
    for i in range(1, N + 1):
        for j in range(N + 1):
            prefix_sum[i][j] += prefix_sum[i - 1][j]
    
    queries = [tuple(map(int, line.split())) for line in it]
    output = []
    for x1, y1, x2, y2 in queries:
        res = (prefix_sum[x2][y2] - prefix_sum[x1 - 1][y2] -
               prefix_sum[x2][y1 - 1] + prefix_sum[x1 - 1][y1 - 1])
        output.append(str(res))
    
    sys.stdout.write("\n".join(output))

if __name__ == '__main__':
    main()