Skip to content

segment tree

Jump to bottom
Changi Cho (tube) edited this page Oct 7, 2025 · 14 revisions

세그먼트트리

구간의 값들을 미리 계한새놓고, 이분탐색으로 구간의 연산 결과를 빨리 찾아내는 방법

트리의 크기를 만들 때에는 원본 배열의 크기에 4배로 만들어주면됨

참고자료

새로운 구현

일반 세그먼트 트리

특정 index 값을 업데이트

struct SegTree {
  const int NULL_VALUE = 0;
  int size;
  vector<int> tree;

  SegTree(vector<int> arr) {
    size = arr.size();
    tree.resize(size * 4);
    build(1, 0, size - 1, arr);
  }

  // initialize part
  void build(int node, int start, int end, vector<int>& arr) {
    if (start == end) {
      tree[node] = arr[start];
      return;
    }

    int mid = (start + end) / 2;
    build(node * 2, start, mid, arr);
    build(node * 2 + 1, mid + 1, end, arr);

    tree[node] = tree[node * 2] + tree[node * 2 + 1];
  }

  // query part
  int query(int left, int right) { return _query(1, 0, size - 1, left, right); }

  int _query(int node, int l, int r, int ql, int qr) {
    if (qr < l || r < ql) {
      return NULL_VALUE;
    }
    if (ql <= l && r <= qr) {
      return tree[node];
    }

    int mid = (l + r) / 2;
    int leftPart = _query(node * 2, l, mid, ql, qr);
    int rightPart = _query(node * 2 + 1, mid + 1, r, ql, qr);
    return leftPart + rightPart;
  }

  // update part
  void update(int i, int diff) { _update(1, 0, size - 1, i, i, diff); }

  void _update(int node, int l, int r, int ql, int qr, int diff) {
    if (l > r || qr < l || r < ql) {
      return;
    }
    if (l == r) {
      tree[node] += diff;
      return;
    }

    int mid = (l + r) / 2;
    _update(node * 2, l, mid, ql, qr, diff);
    _update(node * 2 + 1, mid + 1, r, ql, qr, diff);

    tree[node] = tree[node * 2] + tree[node * 2 + 1];
  }
};

Lazy Propagation

늦은 전파를 사용해 구간 변화 업데이트

struct LazySegTree {
  const int NULL_VALUE = 0;
  int size;
  vector<int> tree;
  vector<int> lazy;

  LazySegTree(vector<int> arr) {
    size = arr.size();
    tree.resize(4 * size);
    lazy.resize(4 * size);

    build(arr, 1, 0, size - 1);
  }

  // initialize part
  void build(vector<int>& arr, int node, int l, int r) {
    if (l == r) {
      tree[node] = arr[l];
      return;
    }

    int mid = (l + r) / 2;
    build(arr, node * 2, l, mid);
    build(arr, node * 2 + 1, mid + 1, r);

    tree[node] = tree[node * 2] + tree[node * 2 + 1];
  }

  void propagate(int node, int l, int r) {
    if (lazy[node] == 0) return;

    tree[node] += lazy[node];
    if (l != r) {
      lazy[2 * node] += lazy[node];
      lazy[2 * node + 1] += lazy[node];
    }
    lazy[node] = 0;
  }

  // update part
  void update(int left, int right, int val) {
    _update(1, 0, size - 1, left, right, val);
  }

  void _update(int node, int l, int r, int ql, int qr, int val) {
    propagate(node, l, r);

    if (l > r || l > qr || r < ql) return;
    if (ql <= l && r <= qr) {
      lazy[node] += val;
      propagate(node, l, r);
      return;
    }

    int mid = (l + r) / 2;
    _update(node * 2, l, mid, ql, qr, val);
    _update(node * 2 + 1, mid + 1, r, ql, qr, val);

    tree[node] = tree[node * 2] + tree[node * 2 + 1];
  }

  // query part
  int query(int left, int right) { return _query(1, 0, size - 1, left, right); }

  int _query(int node, int l, int r, int ql, int qr) {
    propagate(node, l, r);
    if (l > r || qr < l || ql > r) return NULL_VALUE;
    if (ql <= l && r <= qr) return tree[node];

    int mid = (l + r) / 2;
    return _query(node * 2, l, mid, ql, qr) +
           _query(node * 2 + 1, mid + 1, r, ql, qr);
  }
};
이전 구현 방법 (lazy propagation 부분 오류)

구현

일반 세그먼트 트리 (index update, range update)

template <class T>
struct SegmentTree {
  const T NULL_VALUE = 0;
  int size;
  vector<T> tree;

  SegmentTree(vector<T> &array) {
    size = array.size();

    tree.resize(size * 4);

    init(1, 0, size - 1, array);
  }

  // operation part

  T operation(T x, T y) { return x + y; }

  // initialize part

  void init(int node, int start, int end, vector<T> &array) {
    if (start == end) {
      tree[node] = array[start];
      return;
    }

    int mid = (start + end) / 2;
    init(node * 2, start, mid, array);
    init(node * 2 + 1, mid + 1, end, array);

    tree[node] = operation(tree[node * 2], tree[node * 2 + 1]);
  }

  // query part

  T query(int left, int right) { return query(1, 0, size - 1, left, right); }

  T query(int node, int start, int end, int left, int right) {
    if (right < start || end < left) {
      return NULL_VALUE;
    }
    if (left <= start && end <= right) {
      return tree[node];
    }

    int mid = (start + end) / 2;
    T leftPart = query(node * 2, start, mid, left, right);
    T rightPart = query(node * 2 + 1, mid + 1, end, left, right);
    return operation(leftPart, rightPart);
  }

  // update part

  void update(int index, T diff) { update(1, 0, size - 1, index, index, diff); }

  void update(int left, int right, T diff) {
    update(1, 0, size - 1, left, right, diff);
  }

  void update(int node, int start, int end, int left, int right, T diff) {
    if (right < start || end < left) {
      return;
    }
    if (start == end) {
      tree[node] += diff;
      return;
    }

    int mid = (start + end) / 2;
    update(node * 2, start, mid, left, right, diff);
    update(node * 2 + 1, mid + 1, end, left, right, diff);

    tree[node] = operation(tree[node * 2], tree[node * 2 + 1]);
  }
};

lazy propagation segment tree

lazy propagation을 이용해 range를 update하는 메소드 추가

template <class T>
struct SegmentTree {
  const T NULL_VALUE = 0;
  int size;
  vector<T> tree;
  vector<T> lazyDiff;

  SegmentTree(vector<T> &array) {
    size = array.size();

    tree.resize(size * 4);
    lazyDiff.resize(size * 4);

    init(1, 0, size - 1, array);
  }

  // operation part

  T operation(T x, T y) { return x + y; }

  // initialize part

  void init(int node, int start, int end, vector<T> &array) {
    if (start == end) {
      tree[node] = array[start];
      return;
    }

    int mid = (start + end) / 2;
    init(node * 2, start, mid, array);
    init(node * 2 + 1, mid + 1, end, array);

    tree[node] = operation(tree[node * 2], tree[node * 2 + 1]);
  }

  // query part

  T query(int left, int right) { return query(1, 0, size - 1, left, right); }

  T query(int node, int start, int end, int left, int right) {
    // only use lazy propagation
    _lazy(node, start, end);

    if (right < start || end < left) {
      return NULL_VALUE;
    }
    if (left <= start && end <= right) {
      return tree[node];
    }

    int mid = (start + end) / 2;
    T leftPart = query(node * 2, start, mid, left, right);
    T rightPart = query(node * 2 + 1, mid + 1, end, left, right);
    return operation(leftPart, rightPart);
  }

  // update part

  void update(int index, T diff) { update(1, 0, size - 1, index, index, diff); }

  void update(int left, int right, T diff) {
    update(1, 0, size - 1, left, right, diff);
  }

  void update(int node, int start, int end, int left, int right, T diff) {
    if (right < start || end < left) {
      return;
    }
    if (start == end) {
      tree[node] += diff;
      return;
    }

    int mid = (start + end) / 2;
    update(node * 2, start, mid, left, right, diff);
    update(node * 2 + 1, mid + 1, end, left, right, diff);

    tree[node] = operation(tree[node * 2], tree[node * 2 + 1]);
  }

  // update lazy part

  void _lazy(int node, int start, int end) {
    if (lazyDiff[node] == 0) {
      return;
    }

    tree[node] += (end - start + 1) * lazyDiff[node];
    if (start != end) {
      lazyDiff[node * 2] += lazyDiff[node];
      lazyDiff[node * 2 + 1] += lazyDiff[node];
    }
    lazyDiff[node] = 0;
  }

  void update_lazy(int left, int right, T diff) {
    update_lazy(1, 0, size - 1, left, right, diff);
  }

  void update_lazy(int node, int start, int end, int left, int right, T diff) {
    _lazy(node, start, end);

    if (right < start || end < left) {
      return;
    }
    if (start == end) {
      tree[node] += diff;
      return;
    }
    if (left <= start && end <= right) {
      tree[node] += (end - start + 1) * diff;

      lazyDiff[node * 2] += diff;
      lazyDiff[node * 2 + 1] += diff;
      return;
    }

    int mid = (start + end) / 2;
    update_lazy(node * 2, start, mid, left, right, diff);
    update_lazy(node * 2 + 1, mid + 1, end, left, right, diff);

    tree[node] = operation(tree[node * 2], tree[node * 2 + 1]);
  }
};

Clone this wiki locally