2차원 행렬 matrix가 주어졌을 때, 왼쪽 위 모서리(row1, col1)와 오른쪽 아래 모서리(row2, col2)로 정의되는 직사각형 영역 안에 있는 모든 요소의 합을 계산하는 문제입니다. 여기에 더해, 행렬의 특정 위치 값을 갱신(update)한 후에도 같은 연산을 효율적으로 수행할 수 있어야 합니다.
예를 들어 입력이 다음과 같다고 가정해 보겠습니다.
| 3 | 0 | 1 | 4 | 2 |
| 5 | 6 | 3 | 2 | 1 |
| 1 | 2 | 0 | 1 | 5 |
| 4 | 1 | 0 | 1 | 7 |
| 1 | 0 | 3 | 0 | 5 |
이 행렬에 대해 다음과 같이 메서드를 호출한다고 가정합니다.
sumRegion(2, 1, 4, 3)
update(3, 2, 2)
sumRegion(2, 1, 4, 3)
그러면 출력은 8과 10이 됩니다. 초록색으로 표시된 직사각형은 (2,1)과 (4,3) 두 좌표로 정의되며, 처음 이 영역의 합은 8입니다. 이후 update(3, 2, 2)로 값이 갱신되면 같은 영역의 합은 10으로 변합니다.
해결 접근 방식
매번 영역을 직접 순회하며 합을 구하면 한 번의 쿼리에 O(n·m)의 시간이 걸립니다. 조회와 갱신이 반복되는 상황에서는 2차원 펜윅 트리(Binary Indexed Tree, BIT)를 사용하는 것이 효율적이며, 각 연산을 O(log n · log m)에 처리할 수 있습니다.
구현 절차는 다음과 같습니다.
- 2차원 배열 tree와 value를 정의합니다.
- 행렬을 입력받는 생성자(initializer)를 정의합니다.
- n := 행렬의 행 크기
- m := n이 0이 아니면 행렬의 열 크기, 그렇지 않으면 0
- value := n × m 크기의 2차원 배열로 초기화
- tree := (n + 1) × (m + 1) 크기의 2차원 배열로 초기화
- 모든 i, j에 대해 update(i, j, matrix[i][j])를 호출하여 트리를 구성합니다.
update(row, col, val)
- n 또는 m이 0이면 즉시 반환합니다.
- delta := val - value[row][col]
- value[row][col] := val
- i를 row + 1부터 시작해 i ≤ n인 동안 i += i & (-i)로 갱신하며 반복하고, 내부에서는 j를 col + 1부터 j ≤ m까지 j += j & (-j)로 갱신하며 tree[i][j] += delta를 수행합니다.
sum(row, col)
- ret := 0으로 초기화합니다.
- i를 row부터 시작해 i > 0인 동안 i -= i & (-i)로 갱신하며 반복하고, 내부에서는 j를 col부터 j > 0까지 j -= j & (-j)로 갱신하며 ret += tree[i][j]를 누적합니다.
- ret을 반환합니다.
sumRegion(row1, col1, row2, col2)
- m 또는 n이 0이면 0을 반환합니다.
- row2, row1, col1, col2를 각각 1씩 증가시킵니다. (펜윅 트리는 1-based 인덱스를 사용하기 때문입니다.)
- 포함-배제 원리에 따라 sum(row2, col2) + sum(row1 - 1, col1 - 1) - sum(row1 - 1, col2) - sum(row2, col1 - 1)을 반환합니다.
예제 코드 (C++)
더 나은 이해를 위해 다음 구현 예시를 살펴보겠습니다.
#include <bits/stdc++.h>
using namespace std;
class NumMatrix {
public:
int n, m;
vector<vector<int>> tree;
vector<vector<int>> value;
NumMatrix(vector<vector<int>> &matrix) {
n = matrix.size();
m = !n ? 0 : matrix[0].size();
value = vector<vector<int>>(n, vector<int>(m));
tree = vector<vector<int>>(n + 1, vector<int>(m + 1));
for (int i = 0; i < n; i++) {
for (int j = 0; j < m; j++) {
update(i, j, matrix[i][j]);
}
}
}
void update(int row, int col, int val) {
if (n == 0 || m == 0)
return;
int delta = val - value[row][col];
value[row][col] = val;
for (int i = row + 1; i <= n; i += i & (-i)) {
for (int j = col + 1; j <= m; j += j & (-j)) {
tree[i][j] += delta;
}
}
}
int sum(int row, int col) {
int ret = 0;
for (int i = row; i > 0; i -= i & (-i)) {
for (int j = col; j > 0; j -= j & (-j)) {
ret += tree[i][j];
}
}
return ret;
}
int sumRegion(int row1, int col1, int row2, int col2) {
if (m == 0 || n == 0)
return 0;
row2++;
row1++;
col1++;
col2++;
return sum(row2, col2) + sum(row1 - 1, col1 - 1) - sum(row1 - 1, col2) - sum(row2, col1 - 1);
}
};
main() {
vector<vector<int>> v = {
{3, 0, 1, 4, 2},
{5, 6, 3, 2, 1},
{1, 2, 0, 1, 5},
{4, 1, 0, 1, 7},
{1, 0, 3, 0, 5}};
NumMatrix ob(v);
cout << (ob.sumRegion(2, 1, 4, 3)) << endl;
ob.update(3, 2, 2);
cout << (ob.sumRegion(2, 1, 4, 3)) << endl;
}
입력
vector<vector<int>> v = {
{3, 0, 1, 4, 2},
{5, 6, 3, 2, 1},
{1, 2, 0, 1, 5},
{4, 1, 0, 1, 7},
{1, 0, 3, 0, 5}};
NumMatrix ob(v);
cout << (ob.sumRegion(2, 1, 4, 3)) << endl;
ob.update(3, 2, 2);
cout << (ob.sumRegion(2, 1, 4, 3)) << endl;출력
8
10
시간 복잡도 분석
- 생성자: 모든 칸에 대해 update를 호출하므로 O(n · m · log n · log m)
- update(): O(log n · log m)
- sumRegion(): 네 번의 sum 호출로 구성되므로 O(log n · log m)
이처럼 2차원 펜윅 트리를 활용하면 행렬 값이 자주 바뀌는 환경에서도 빠른 영역 합 조회가 가능합니다.