백준 2167번: 2차원 배열의 합 (S5) - Python 풀이

문제 설명

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

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가 주어진다(1 ≤ i ≤ x ≤ N, 1 ≤ j ≤ y ≤ M).

출력

K개의 줄에 순서대로 배열의 합을 출력한다. 배열의 합은 231-1보다 작거나 같다.

예제 입력 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

문제 설명이 좀 불친절합니다.

문제 요약) (i,j)에서 (x,y)까지 네모 그리고 그 네모 안의 모든 숫자의 합을 출력해라~

https://www.acmicpc.net/board/view/1315#comment-4205

풀이 과정

[첫 번째 시도] 처음에는 K만큼 도는 for문 안에서, i부터 x까지의 row, j부터 y까지의 col 값을 2개의 for문으로 각각 구해서 다 더하는 식으로 풀었습니다. → 시간 초과!

⇒ for문 3개 중첩이라 그런가? for문 개수를 줄이면 되려나 생각했습니다.

[두 번째 시도]

그래서 주어진 2차원 배열에서 (i,j) ~ (x,y) 범위만큼 슬라이싱 해서 가져오는 것을 생각했습니다.

2차원 배열을 슬라이싱 하려면, 1차원 배열을 슬라이싱 하는 것처럼 한 번에 하는 방법은 없습니다. 원하는 행만큼 돌면서 열 기준으로 슬라이싱한 값(= 1차원 배열 슬라이싱)을 더해서 구해야 합니다.

슬라이싱 한 값들을 1차원 배열에 저장해서 sum으로 총합을 구하도록 풀었습니다.

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

for i in range(N):
    arr[i] += list(map(int, input().split()))

K = int(input())
for _ in range(K):
    i, j, x, y = map(int, input().split())
    tmp = []
    for row in range(i-1, x):
        tmp += arr[row][j-1:y] # 2차원 배열 슬라이싱
    print(sum(tmp))

→ 시간 초과!

for문 2개면 괜찮지 않나 했는데 ㅠㅠ

⇒ 그 뒤로 질문 게시판을 보니 누적합을 사용해야 하는 문제였습니다.

누적 합 (Prefix Sum)

• 구간의 합을 빠르게 계산해야 할 때 사용하는 알고리즘
• 주어진 배열의 앞에서부터 차례대로 누적된 합을 구해놓고, 이를 이용해서 구간의 합을 구할 수 있습니다.
• 배열 A의 누적 합 배열 S는 다음과 같이 정의됩니다.
  • 이때, l번째 수에서 r번째 수까지의 합은 S[r] - S[l-1] 과 같습니다.S[l-1] = A[0] + … + A[l-1]
  • 이기 때문에, S[r] - S[l-1] = A[l] + … A[r]이 됩니다.
  • S[r] = A[0] + … + A[r]
• S[i] = A[0] + A[1] + … + A[i]
• 시간복잡도
  • 누적 합 배열 생성 시: O(n)
  • 구간 합 계산 시: O(1)

이 문제도 특정 구간의 합을 빠르게 구해야 하기 때문에 누적 합을 사용합니다.

정답 코드

N, M = map(int, input().split())
arr = [[] for _ in range(N)]
# 누적 합을 저장할 2차원 배열
sum_arr = [[0 for _ in range(M)] for _ in range(N)]

for i in range(N):
    l = list(map(int, input().split()))
    arr[i] += l
    # 각 행의 첫 번째 값은 입력의 첫 번째 값을 그대로
    sum_arr[i][0] = l[0]
    # 그 이후의 열은 앞에 값과 누적해서 계산
    for j in range(1, M):
        # 바로 앞까지 누적된 값에 l[i]를 더한다
        sum_arr[i][j] += sum_arr[i][j-1] + l[j]

K = int(input())
for _ in range(K):
    i, j, x, y = map(int, input().split())
    ans = 0
    # 행마다 누적 합을 구해놨기 때문에 (i-1)부터 (x-1)까지의 행을 확인
    for row in range(i-1, x):
        # 누적 합으로 가져와야 하는 구간: (열의 최소 인덱스) ~ (열의 최대 인덱스)
        # 열의 최댓값인 y-1에서
        ans += sum_arr[row][y-1]
        # j가 1 이하이면 구간의 맨 처음이라서 뺄 것이 없기 때문에 조건으로 걸러준다
        if j > 1:
            # 열의 최솟값인 j-1의 직전까지 즉 j-2까지의 누적 합을 뺀다
            # => y-1열에서 j-1열까지의 구간 합을 구할 수 있다 
            ans -= sum_arr[row][j-2]
    print(ans)

[세 번째 시도] 행마다 누적 합을 구해서 sum_arr에 넣어두고, 범위가 주어지면 행마다 가져와야 하는 열의 구간에 맞춰서 구해둔 누적 합을 이용해 빠르게 계산한다. → 정답!

이상입니다.