[누적합] 세그먼트트리 / Segment tree

세그먼트 트리에 대해서 알아보고 해당하는 자료구조에 대해서 코드로 어떻게 사용하는지에 대해서 정리하려고 작성하였습니다.

https://book.acmicpc.net/ds/segment-tree

자세한 설명은 위의 백준북에서 볼 수 있습니다.

 

 

C++

#include <bits/stdc++.h>

using namespace std;
using ll = long long;

ll init(vector<ll> &a, vector<ll> &tree, int node, int start, int end);
ll sum(vector<ll> &tree, int node, int start, int end, int left, int right);
void update(vector<ll> &tree, int node, int start, int end, int index, ll diff);

int main() {
  //// Segment tree
  // 세그먼트 트리에서 노드들은 다음과 같은 의미를 가진다.
  // 리프 노드 : 배열의 그 수 자체.
  // 리프 노드를 제외한 다른 노드 :
  // 왼쪽 자식과 오른쪽 자식의 최솟값 또는 합을저장

  // 트리 배열의 크기. 최적화된 크기는 아래와 같이 구할 수 있다.
  // n이 10인경우에 최적화된 크기는 32
  // 이게 귀찮다면 그냥 n에 4를 곱하면 된다. 모든 경우에 우리가 필요로하는
  // 배열의 크기보다 크기 떄문

  vector<ll> n_arr = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};
  int n = n_arr.size();
  int h = (int)ceil(log2(n));
  int tree_size = (1 << (h + 1));
  vector<ll> tree(tree_size);

  // 세그먼트 트리 초기화. 초기화과정에서의 시간 복잡도는 O(N)
  init(n_arr, tree, 1, 0, n - 1);

  // 합을 구하는 과정에서 시간 복잡도는 O(log N)
  // 합은 4가지로 구분된다.
  // 1.하나의 노드로 원하는 구역의 합을 구하는 경우 예) 0~9까지의 합 (루트노드)
  // 2.left를 리프로 갖고 right를 구역 노드로 갖는 경우 (2, 3~4)
  // 3.right를 리프로 갖고 left를 구역 노드로 갖는 경우 (5~7, 8)
  // 1~3을 제외한 나머지 left, right 전부 구역 노드로 갖는경우 (3~4, 5~9)
  cout << "Sum from 0 to 9: " << sum(tree, 1, 0, n - 1, 0, 9) << "\n";
  cout << "Sum from 2 to 4: " << sum(tree, 1, 0, n - 1, 2, 4) << "\n";
  cout << "Sum from 5 to 8: " << sum(tree, 1, 0, n - 1, 5, 8) << "\n";
  cout << "Sum from 3 to 9: " << sum(tree, 1, 0, n - 1, 3, 9) << "\n";

  // 업데이트 과정에서 시간 복잡도는 O(log N)
  // 인덱스 3의 값을 7로 변경 (기존 값은 3)
  int index = 3;
  ll new_value = 7;
  ll diff = new_value - n_arr[index];
  n_arr[index] = new_value;
  update(tree, 1, 0, n - 1, index, diff);  // 세그먼트 트리 업데이트

  // 업데이트 후 구간 합 쿼리
  cout << "After updating index 3 to 7, sum from 0 to 9: "
       << sum(tree, 1, 0, n - 1, 0, 9) << "\n";
  cout << "After updating index 3 to 7, sum from 2 to 4: "
       << sum(tree, 1, 0, n - 1, 2, 4) << "\n";

  // 인덱스 5의 값을 10으로 변경 (기존 값은 5)
  index = 5;
  new_value = 10;
  diff = new_value - n_arr[index];
  n_arr[index] = new_value;                // 배열 업데이트
  update(tree, 1, 0, n - 1, index, diff);  // 세그먼트 트리 업데이트

  // 업데이트 후 구간 합 쿼리
  cout << "After updating index 5 to 10, sum from 0 to 9: "
       << sum(tree, 1, 0, n - 1, 0, 9) << "\n";
  cout << "After updating index 5 to 10, sum from 5 to 8: "
       << sum(tree, 1, 0, n - 1, 5, 8) << "\n";

  return 0;
}

ll init(vector<ll> &a, vector<ll> &tree, int node, int start, int end) {
  // a: 초기 배열, tree: 세그먼트 트리, node: 세그먼트 트리 노드 번호
  // node가 담당하는 합의 범위가 start ~ end
  if (start == end)                // 노드가 리프 노드인 경우
    return tree[node] = a[start];  // 배열의 그 원소를 가져야 함

  int mid = (start + end) / 2;

  // 구간 합을 구하는 경우
  return tree[node] = init(a, tree, node * 2, start, mid) +
                      init(a, tree, node * 2 + 1, mid + 1, end);
}

ll sum(vector<ll> &tree, int node, int start, int end, int left, int right) {
  // case 1: [start, end] 앞 뒤에 [left, right]가 있는 경우,
  // 겹치지 않기 때문에 탐색을 더 이상 할 필요가 없다.
  if (left > end || right < start) return 0;

  // case 2: [start, end]가 [left, right]에 포함
  if (left <= start && end <= right) return tree[node];

  // case 3, 4: 왼쪽 자식과 오른쪽 자식을 루트로 하는 트리에서 다시 탐색 시작
  int mid = (start + end) / 2;
  return sum(tree, node * 2, start, mid, left, right) +
         sum(tree, node * 2 + 1, mid + 1, end, left, right);
}

void update(vector<ll> &tree, int node, int start, int end, int index,
            ll diff) {
  if (index < start || index > end) return;  // case 2
  tree[node] = tree[node] + diff;            // case 1

  // 리프 노드가 아닌 경우 자식도 변경해줘야 하기 때문에,
  // 리프 노드인지 검사를 하고 아래 자식 노드를 갱신해준다.
  if (start != end) {
    int mid = (start + end) / 2;
    update(tree, node * 2, start, mid, index, diff);
    update(tree, node * 2 + 1, mid + 1, end, index, diff);
  }
}