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

문제

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

 

 

풀이

무식하게 그냥 더하거나 (PyPy3만 통과 가능), 누적합 배열을 사용하거나 둘 다 가능하다 

 

 

누적합을 이용하는법의 풀이는 다음과 같다.

먼저 보라색 영역의 합을 구한다고 가정하자

정답은 6+7+10+11 = 34이다

 

색칠된 모든 영역을 더한 다음 (1+2+3 + 5+6+7 + 9+10+11 = 54)

초록색 영역 (1+2+3 = 6)을 빼고

파란색 영역인 (1+5+9 = 15)를 빼고

겹쳐서 두번 빼버린 1을 다시 더하자

결과는 54 - 6 - 15 + 1 = 34로 보라색 영역만 더한 값이 나온다.

 

미리 "(0,0)부터 각각의 좌표까지의 영역"을 더해서 저장해놓은 누적합 이차원 배열을 만들고

보라색-초록색-파란색+빨간색 위치의 숫자만 계산하면 빠르게 답을 구할 수 있다.

(초록색, 파란색, 빨간색 영역의 인덱스가 0보다 작아지는 경우의 예외 처리를 주의할것)

 

 

파이썬에서 1차원 리스트의 누적합은 itertools 의 accumulate로 쉽게 구할 수 있고

리스트와 리스트의 원소간 덧셈은 operator 의 add를 이용해  list(map(add, 리스트A, 리스트B))와 같이 구할 수 있으며

1차원의 누적합 리스트를 원소끼리 계속해서 더해가면 2차원 누적합을 구할 수 있다.
(numpy로 한방에 2차원 누적합을 구할 수 있지만... 코테에선 보통 못쓴다)

 

 

 

 

정답 코드 (무식하게 더하기)

import sys
input = sys.stdin.readline

N, M = map(int, input().split())
numArray = [list(map(int, input().split())) for _ in range(N)]

K = int(input())
for _ in range(K):
    answer = 0
    i,j,x,y = map(lambda n: n-1, map(int, input().split()))

    for row in range(i, x+1):
        answer += sum(numArray[row][j:y+1])

    print(answer)

 

정답 코드 (누적합)

from itertools import accumulate
from operator import add
import sys
input = sys.stdin.readline

N, M = map(int, input().split())
rowPrefixSum = [list(accumulate(list(map(int, input().split())))) for _ in range(N)]
prefixSum = []
for i in range(N):
    if i==0:
        prefixSum.append(rowPrefixSum[i])
        continue
    else:
        prefixSum.append(list(map(add, prefixSum[i-1], rowPrefixSum[i])))

T = int(input())
for _ in range(T):
    i,j,x,y = map(lambda n : n - 1, map(int, input().split()))
    answer = prefixSum[x][y]

    if i>0:
        answer -= prefixSum[i-1][y]
    if j>0:
        answer -= prefixSum[x][j-1]
    if i>0 and j>0:
        answer += prefixSum[i-1][j-1]
    print(answer)