[백준2167번] 2차원 배열의 합 / Python3

문제

2차원 배열이 주어졌을 때 (i, j) 위치부터 (x, y) 위치까지에 저장되어 있는 수들의 합을 구하는 프로그램을 작성하시오. 배열의 (i, j) 위치는 i행 j열을 나타낸다.

입력

첫째 줄에 배열의 크기 N, M(1 ≤ N, M ≤ 300)이 주어진다. 다음 N개의 줄에는 M개의 정수로 배열이 주어진다. 배열에 포함되어 있는 수는 절댓값이 10,000보다 작거나 같은 정수이다. 그 다음 줄에는 합을 구할 부분의 개수 K(1 ≤ K ≤ 10,000)가 주어진다. 다음 K개의 줄에는 네 개의 정수로 i, j, x, y가 주어진다(i ≤ x, j ≤ y).

출력

K개의 줄에 순서대로 배열의 합을 출력한다. 배열의 합은 32bit-int 범위를 초과하지 않는다.

예제 입력 1

2 3
1 2 4
8 16 32
3
1 1 2 3
1 2 1 2
1 3 2 3

예제 출력 1

63
2
36

나의 풀이

이 문제는 얼핏 보면 그냥 k번만큼 이중for문을 돌면서 모두 더해주면 쉽게 풀릴 것이라고 생각할 수 있다. 그러나 k의 범위는 10000까지로 작은 수가 아니다. 앞선 생각으로 풀게되면 아마도 시간초과가 뜰것이다.

필자는 이 문제를 최초 한번 dp에 각 행에 대하여 합들을 구해놓고 k번 합을 구할때에는 for문을 한번씩만 돌게끔 만들어 주었다. 우선 위의 예제에 대해 행에 대한 dp를 모두 구해보면 다음과 같을 것이다.

1 3 7
8 24 56

단순히 열이 증가할수록 이전 열의 값에 더해주는 것이다.

이렇게 행마다 합에 대한 dp를 만들어 놓으면 어떻게 이중포문이 아닌 한번의 for문으로 문제를 풀 수 있는 것일까? 답은 쉽다. 위에서 1행에 대한 경우만 생각해보자. 가령 2부터 3까지의 합을 구하고자 한다면 2+4 == 6 == 7-1이다. 이것은 다음과 같다. dp[3] - dp[1] (1행만을 생각) 따라서 각 행에 대한 원하는 구간의 합을 구하는 점화식은 다음과 같다.

// r행에서 j부터 y까지의 합
result = dp[r][y] - dp[r][j-1]

이렇게 점화식까지 알았다면 행의 값인 r을 주어진는 행의 값 i부터 x까지 for문을 돌면서 모든 행에서 구해지는 값을 더해 출력해주면 정답이다.

코드

# 2167번 2차원 배열의 합
import sys

# main
n, m = [int(x) for x in sys.stdin.readline().split()]
arr = [[int(x) for x in sys.stdin.readline().split()] for _ in range(n)]

# 합에 대한 dp 생성
dp = [[0 for _ in range(m+1)] for _ in range(n+1)]
for i in range(1,n+1):
    for j in range(1,m+1):
        dp[i][j] = dp[i][j-1] + arr[i-1][j-1]

# 범위 합 구하기
k = int(sys.stdin.readline())
for _ in range(k): 
    i, j, x, y = [int(x) for x in sys.stdin.readline().split()]

    result = 0
    for r in range(i,x+1):
        result += dp[r][y] - dp[r][j-1]

    print (result)