Computer >> 컴퓨터 >  >> 프로그래밍 >> C++

C++로 해결하는 2D 범위 합 쿼리(변경 가능) – 2차원 펜윅 트리 활용법

2차원 행렬 matrix가 주어졌을 때, 왼쪽 위 모서리(row1, col1)와 오른쪽 아래 모서리(row2, col2)로 정의되는 직사각형 영역 안에 있는 모든 요소의 합을 계산하는 문제입니다. 여기에 더해, 행렬의 특정 위치 값을 갱신(update)한 후에도 같은 연산을 효율적으로 수행할 수 있어야 합니다.

예를 들어 입력이 다음과 같다고 가정해 보겠습니다.

30142
56321
12015
41017
10305

이 행렬에 대해 다음과 같이 메서드를 호출한다고 가정합니다.

sumRegion(2, 1, 4, 3)
update(3, 2, 2)
sumRegion(2, 1, 4, 3)

그러면 출력은 810이 됩니다. 초록색으로 표시된 직사각형은 (2,1)과 (4,3) 두 좌표로 정의되며, 처음 이 영역의 합은 8입니다. 이후 update(3, 2, 2)로 값이 갱신되면 같은 영역의 합은 10으로 변합니다.

해결 접근 방식

매번 영역을 직접 순회하며 합을 구하면 한 번의 쿼리에 O(n·m)의 시간이 걸립니다. 조회와 갱신이 반복되는 상황에서는 2차원 펜윅 트리(Binary Indexed Tree, BIT)를 사용하는 것이 효율적이며, 각 연산을 O(log n · log m)에 처리할 수 있습니다.

구현 절차는 다음과 같습니다.

  • 2차원 배열 treevalue를 정의합니다.
  • 행렬을 입력받는 생성자(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차원 펜윅 트리를 활용하면 행렬 값이 자주 바뀌는 환경에서도 빠른 영역 합 조회가 가능합니다.