[PS] BOJ 11660 / 구간 합 구하기 5

[PS] BOJ 11660 / 구간 합 구하기 5
문제 링크: https://www.acmicpc.net/problem/11660
Thumbnail: Photo by Melissa (Unsplash)

2차원 배열의 누적 합 문제입니다! 비슷한 문제를 풀었던 기억이 있네요.

풀이

누적 합을 계산하는 풀이는 지정좌석제 (33993)와 같습니다.

다만, 지정좌석제 풀이에서는 범위의 중앙 좌표를 기준으로 계산했다면 이번 문제는 편의 상 끝 쪽 좌표를 사용했습니다.

누적 합 구하기

먼저, 입력으로 주어진 2차원 배열의 누적 합을 계산해야 합니다. 도형의 넓이를 계산한다고 생각하면 쉽게 식을 떠올릴 수 있습니다.

N, M = map(int, input().strip().split())
arr = [list(map(int, input().strip().split())) for _ in range(N)] # 입력 배열 (N x N)
prefix_sum = [[0 for _ in range(N)] for _ in range(N)] # 누적 합 배열

for i in range(N):
    for j in range(N):
        if i == 0:
            if j == 0:
                prefix_sum[i][j] = arr[0][0]
            else:
                prefix_sum[i][j] = arr[i][j] + prefix_sum[i][j - 1]
        elif j == 0:
            prefix_sum[i][j] = arr[i][j] + prefix_sum[i - 1][j]
        else:
            prefix_sum[i][j] = arr[i][j] + prefix_sum[i-1][j] + prefix_sum[i][j - 1] - prefix_sum[i - 1][j - 1]

임의의 구간 합 구하기

앞서 계산한 누적 합으로 임의의 구간 합을 계산하는 방법을 알아봅시다. 도형의 넓이를 계산한다고 생각하면 이해하기 쉽습니다.

2차원 배열의 구간 합 계산 방식
# 구간별 합 출력하기
for _ in range(M):
    x1, y1, x2, y2 = map(lambda x: int(x) - 1, input().strip().split()) # 실제 인덱스는 0부터 시작하므로 1 줄인다.
    
    # 경우의 수 처리
    if x1 == 0:
        if y1 == 0: # 추가 계산이 필요 없는 경우.
            print(prefix_sum[x2][y2])
        else:  # y1이 0이 아닌 경우
            print(prefix_sum[x2][y2] - prefix_sum[x2][y1 - 1])
    elif y1 == 0:  # x1이 0이 아닌 경우
        print(prefix_sum[x2][y2] - prefix_sum[x1 - 1][y2])
    else:   # 일반적인 경우
        print(prefix_sum[x2][y2] - prefix_sum[x1 - 1][y2] - prefix_sum[x2][y1 - 1] + prefix_sum[x1 - 1][y1 - 1])

구간 합 계산하기

전체 코드

input = open(0).readline

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

# 누적 합 계산하기
for i in range(N):
    for j in range(N):
        if i == 0:
            if j == 0:
                prefix_sum[i][j] = arr[0][0]
            else:
                prefix_sum[i][j] = arr[i][j] + prefix_sum[i][j - 1]
        elif j == 0:
            prefix_sum[i][j] = arr[i][j] + prefix_sum[i - 1][j]
        else:
            prefix_sum[i][j] = arr[i][j] + prefix_sum[i-1][j] + prefix_sum[i][j - 1] - prefix_sum[i - 1][j - 1]

# 구간별 합 출력하기
for _ in range(M):
    x1, y1, x2, y2 = map(lambda x: int(x) - 1, input().strip().split()) # 실제 인덱스는 0부터 시작하므로 1 줄인다.
    
    # 경우의 수 처리
    if x1 == 0:
        if y1 == 0: # 추가 계산이 필요 없는 경우.
            print(prefix_sum[x2][y2])
        else:  # y1이 0이 아닌 경우
            print(prefix_sum[x2][y2] - prefix_sum[x2][y1 - 1])
    elif y1 == 0:  # x1이 0이 아닌 경우
        print(prefix_sum[x2][y2] - prefix_sum[x1 - 1][y2])
    else:   # 일반적인 경우
        print(prefix_sum[x2][y2] - prefix_sum[x1 - 1][y2] - prefix_sum[x2][y1 - 1] + prefix_sum[x1 - 1][y1 - 1])

solution.py