[알고리즘] 구간 합 (+LeetCode 1314)

알고리즘 (1차원)

구간 합 (prefix sum)은 숫자 배열이 \( [x_1,x_2,...,x_n] \)이 있을 때 다음과 같이 구간의 합을 구해서 새로운 숫자 배열 \([y_1,y_2,...,y_n]\)을 만들어서 사용하는 방법이다. 여기서 y들은 다음과 같이 정의된다:

• \(y_1=x_1\)
• \(y_2=x_1+x_2\)
• \(y_3=x_1+x_2+x_3\)
• \(y_n=x_1+x_2+x_3+...+x_n\)

실제로 계산할 때는 \(y_0=0\)으로 두고 다음과 같이 계산하는 것이 편하다

• \(y_i = x_i+y_{i-1}, \quad i=1,2...,n\)
• \( y_0=0, y_1=x_1+0=x_1, y_2=y_1+x_2=x_2+x_1,...\)인 것을 확인할 수 있다.

이렇게 계산하는 과정의 계산 복잡도는 \( O(n)\)이 된다.

사용 예시

배열 \( [x_1,x_2,...,x_n] \)이 있고, k개의 구간 \((s_1,e_1), (s_2,e_2),...,(s_n,e_n) \)이 있다고 하자. 여기서 s는 시작점, e는 끝점을 나타내며, s에서 e 사이에 있는 합 (시작점 끝점 포함) 을 구하고 싶다고 하자. 구간의 길이가 O(n)이라고 가정하자. 단순 무식하게 매번 \( x_{s_i} + ... + x_{e_i}\)를 계산하면 계산 복잡도는 O(n x k)가 된다 (O(n)의 계산을 k번 반복). 하지만 구간 합 y를 계산해 두면  s에서 e 사이에 있는 합은 \( y_{e_i}-y_{s_i-1}\)로 계산할 수 있어서 계산 복잡도는 O(n)+O(k)가 된다 (구간 합 계산할 때 O(n), k번 계산할 때 O(k)).

알고리즘 (2차원)

고차원 숫자 배열에서도 비슷한 알고리즘을 사용할 수 있는데 편의상 2차원에 집중하겠다. 2차원의 숫자 배열 \( x_{i,j}, i=1,...,m, j=1,...,n\)이 있다고 하자. 구간 합 \( y_{r,c}\)는 \( i<=r, j<=c\)를 만족하는 \( x_{i,j}\)를 더한 값으로 정의한다.

실제로 다음과 같이 계산하는 것이 편하다

• \( y_{i,j}=y_{i-1,j}+y_{i,j-1}-y_{i-1,j-1}+x_{i,j}\)
• 이전처럼 \( y_{0,j} = 0\), \( y_{i,0} = 0\)으로 두고 계산한다

그림으로 나타내면 다음과 같다

y_33=y_23+y_32-y_22+x_33을 나타낸 그림. 회색으로 칠한 y의 값 (y33=45)를 구하는 것이 목표. x에 해당하는 태이블에서 연두색으로 색칠한 숫자들은 더했다고 생각하면 된다(풀어쓰면 27+21-12+9=45)

1차원과 비슷하게 index 값 \( (s_y,s_x) \) 와 \( (e_y,e_x) \) 사이 \( x\)의 합을 구하고 싶다고 하자 (\( \sum_{i=s_y}^{e_y}  \sum_{j=s_x}^{e_x} x_{ij} \)). 구간 합 y를 구해두면 다음과 같이 구할 수 있다:

\[  \sum_{i=s_y}^{e_y}  \sum_{j=s_x}^{e_x} x_{ij}  = y_{e_y,e_x} - y_{e_y,s_x-1}-y_{s_y-1,e_x}+y_{s_y-1,s_x-1}\]

참고로 이전처럼 \( y_{0,j} = 0\), \( y_{i,0} = 0\)으로 두고 계산한다. 다음 그림을 보면 이해가 빠를 것이다:

\( s_x=2,s_y=2,e_x=3,e_y=3\)을 구하는 과정을 나타낸 그림. 각 테이블이 식 \( y_{e_y,e_x} - y_{e_y,s_x-1}-y_{s_y-1,e_x}+y_{s_y-1,s_x-1}=\sum_{i=s_y}^{e_y}&nbsp; &nbsp;\sum_{j=s_x}^{e_x} x_{ij} &nbsp;\)에 있는 값들에 해당된다

이제 문제를 하나 풀어보자.

https://leetcode.com/problems/matrix-block-sum/

m x n 행렬이 mat이 있을 때 m x n 행렬 answer[i][j]을 구하는 문제다. 여기서 answer[i][j]는  위에 있는 식을 만족하는 r, c에 대해서 mat[i][j]를 더한 값을 갖는다. 위 조건 식을 만족하는  r,c에 대한 합은 (i-k,j-k)와 (i+k,j+k)를 시작점, 끝점으로 갖는 구간 합인 것을 사용하면 풀 수 있다.

class Solution:
    def matrixBlockSum(self, mat: List[List[int]], k: int) -> List[List[int]]:
        h,w = len(mat), len(mat[0]) # mat의 높이와 너비
        # 구간 합(prefix_sum)부터 계산
        prefix_sum = [[0 for j in range(w+1)] for i in range(h+1)] 
        for i in range(h):
            for j in range(w):
                prefix_sum[i+1][j+1] = prefix_sum[i+1][j] + prefix_sum[i][j+1] - prefix_sum[i][j] + mat[i][j]
        # answer 계산
        answer = [[] for i in range(h)]
        for i in range(h):
            for j in range(w):
            	# 구간 합의 높이 방향으로 최소 최대 index. mat의 범위를 벗어나지 않도록 주의
                i1,i2 = max(0,i-k), min(h,i+k+1)
                # 구간 합의 너비 방향으로 최소 최대 index. mat의 범위를 벗어나지 않도록 주의
                j1,j2 = max(0,j-k), min(w,j+k+1) 
                # 구간 합을 사용해서 (i1,j1) & (i2,j2) 사이 mat의 합
                val = prefix_sum[i2][j2]-prefix_sum[i1][j2]-prefix_sum[i2][j1]+prefix_sum[i1][j1] 
                answer[i].append(val)
        return answer

Refs: https://en.wikipedia.org/wiki/Prefix_sum#:~:text=In%20computer%20science%2C%20the%20prefix,0