💻 백준 11660번: 구간 합 구하기 5

문제 소개 🧐

N×N 크기의 표에 채워진 수들이 주어졌을 때, (x1, y1)부터 (x2, y2)까지의 합을 구하는 프로그램을 작성하는 문제다. 합을 구하는 연산이 여러 번 주어지므로 효율적인 방법이 필요했다.

입력 📝

  • 첫째 줄에 표의 크기 N과 합을 구해야 하는 횟수 M이 주어진다. (1 ≤ N ≤ 1024, 1 ≤ M ≤ 100,000)
  • 둘째 줄부터 N개의 줄에는 표에 채워져 있는 수가 1행부터 차례대로 주어진다.
  • 다음 M개의 줄에는 네 개의 정수 x1, y1, x2, y2 가 주어진다.

출력 📤

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

제한 ❌

  • 시간: 1초
  • 메모리: 256 MB

나의 풀이: 누적 합 배열 활용 🚀

Idea

처음에는 각 쿼리마다 이중 반복문으로 합을 구하는 단순한 방법을 생각했다. 하지만 N과 M의 최댓값을 고려했을 때, 1024 * 1024 * 100,000이라는 엄청난 연산 횟수가 예상되어 시간 초과가 날 것이 뻔했다.

그래서 2차원 누적 합(Prefix Sum) 배열을 사용하기로 했다.

  1. 먼저 sum_li[i][j]를 (0, 0)부터 (i, j)까지의 합을 담는 배열로 정의했다.
  2. sum_li 배열을 채울 때, 나는 각 행의 누적 합을 먼저 구하고, 그 값을 이전 행의 같은 열 누적 합과 더하는 방식으로 구현했다.
    • sum_li[i][j] = (현재 행의 0부터 j까지의 합) + sum_li[i-1][j]
  3. 이렇게 미리 계산된 sum_li 배열이 있으면, 특정 구간의 합을 O(1)에 계산할 수 있었다.
    • (x1, y1)부터 (x2, y2)까지의 합은 sum_li 배열의 값들을 더하고 빼는 것으로 간단히 구할 수 있었다.

이 접근법으로 시간 복잡도를 O(N^2 + M)으로 줄여 문제를 해결할 수 있었다.

Code

import sys

n, m = map(int, sys.stdin.readline().split())

num_board = list(list(map(int, sys.stdin.readline().split())) for _ in range(n))

sum_li = [[0 for _ in range(n)] for _ in range(n)]
for i in range(n):
    sum_axis = 0
    for j in range(n):
        sum_axis += num_board[i][j]
        if i == 0:
            sum_li[i][j] = sum_axis
        else:
            sum_li[i][j] = sum_axis  + sum_li[i-1][j]

def print_target_sums(st_x, st_y, end_x, end_y):
    global sum_li
    global num_board

    if st_x > 0 and st_y > 0:
        print(sum_li[end_x][end_y] - sum_li[st_x-1][end_y] - sum_li[end_x][st_y-1] + sum_li[st_x-1][st_y-1])
    elif st_x > 0 and st_y == 0:
        print(sum_li[end_x][end_y] - sum_li[st_x-1][end_y])
    elif st_x == 0 and st_y > 0:
        print(sum_li[end_x][end_y] - sum_li[end_x][st_y-1])
    else:
        print(sum_li[end_x][end_y])
    

for _ in range(m):
    st_x, st_y, end_x, end_y = map(lambda x: int(x) - 1, sys.stdin.readline().split())  # 실제 인덱스에 맞게 -1
    print_target_sums(st_x, st_y, end_x, end_y)

Result

  • 결과: 정답! ✅
  • 누적 합 배열을 미리 계산해 둔 덕분에 각 쿼리를 상수 시간에 처리하여 시간 초과 없이 문제를 해결했다.

개선할 부분 🤔

  • 2차원 배열의 구간 합 문제는 누적 합이 거의 정석적인 풀이법이라는 것을 다시 한번 느꼈다.
  • 내 코드의 누적 합 계산 방식도 문제는 없었지만, sum[i][j] = board[i][j] + sum[i-1][j] + sum[i][j-1] - sum[i-1][j-1] 점화식을 사용하는 것이 더 일반적이고 깔끔한 코드가 될 수 있었을 것 같다.
  • 문제의 좌표는 1-based, 배열 인덱스는 0-based인 점을 항상 주의해야겠다. 이 부분에서 실수가 종종 발생할 수 있을 것 같다.