[코딩 테스트 연습(파이썬/Python)] 백준(BOJ) 11660번 _ 구간 합 구하기 5

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

11660번: 구간 합 구하기 5

[문제]

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)까지 합을 구해 출력한다.

 

[아이디어]

1) 일반적인 for문으로 계산하기에 M의 값이 너무 크다.

2) 누적 합과 관련된 dp를 생성한 뒤, 누적 합의 차로 일정 구간의 합을 구할 수 있다.

 

# Pypy3로 제출 시 통과

n,m = map(int,input().split())
t = [list(map(int,input().split())) for _ in range(n)]
dp = [[0]*n for _ in range(n+1)]
tmp = 0
for i in range(1,n+1):
    for j in range(n):
        tmp += t[i-1][j]
        dp[i][j] = tmp
for _ in range(m):
    a,b,c,d = map(int,input().split())
    tmp = [dp[k][d-1]-dp[k][b-2] if b>=2 else dp[k][d-1]-dp[k-1][-1] for k in range(a,c+1)]
    print(sum(tmp))

각 위치에 누적 합을 저장하면서, 필요한 부분의 합만 구한다. d가 1인 경우 이전 row를 찾아가야 하기 때문에 별도의 조건문이 필요한데, dp를 2차원으로 구성하는게 아니라 1차원으로 구성했다면 불필요한 부분인 것 같다.

 

숏코딩을 살펴보는 과정에서 itertools.accumulate()에 대해서 알게 되었다.

주어진 iterable 객체에 대한 누적 합을 제공하며, initail = n을 통해 시작 값을 지정하거나, sum이 아닌 binary function을 제공함으로써 원하는 연산에 대한 누적 연산의 결과를 받아볼 수 있다. reduce와 비슷한 느낌으로 사용할 수 있다.