Dynamic Segment Tree and Lazy Propagation

이 포스트에서는 다이나믹 세그먼트 트리에 대한 소개와 레이지 프로퍼게이션까지 구현을 해보도록 한다.

엄청 큰 범위에서의 구간 쿼리?

다음 문제를 해결한다고 해보자.

 

수열 a가 있을 때

 

1. ai에 x를 더한다.

 

2. i <= j일 때 ai ~ aj의 합을 구한다.

 

이는 PS 경력이 조금이라도 있다면 세그먼트 트리로 쉽게 해결할 수 있음을 안다.

 

그런데 제한이 다음과 같다고 해보자.

 

모든 원소가 0인 수열 a가 주어진다. 그런데 그 수열의 크기는 최대 $10^{18}$이다.

 

이런 범위가 너무 크다! 비교적 자비로운 문제라면 오프라인 쿼리를 가능하게 한다던가 해서 숨통을 틔워줄 텐데 그마저도 불허한다면?

 

하지만 걱정하지 않아도 된다. 다이나믹 세그먼트 트리는 이 문제를 해결할 수 있게 해 준다.

Dynamic Segment Tree

정적인 노드에서 계산을 하는 세그먼트 트리와 달리 다이나믹 세그먼트 트리는 쿼리를 수행할 때마다 각 범위에 해당하는 노드를 만들게 된다.

 

비어있는 트리에서 10의 17제곱에 해당하는 원소에 5를 더하는 모습

 

그림과 같이 다이나믹 세그먼트 트리는 쿼리를 수행할 때 필요한 노드만 추가하므로 쿼리당 $O(\log N)$의 공간 복잡도가 소비된다.

 

모든 쿼리를 계산할 때 총 공간 복잡도는 $O(Q \log N)$이 된다.

 

노드를 동적으로 잘 추가하면 큰 구간에서의 쿼리는 평범한 세그먼트 트리처럼 풀 수 있다.

 

그러면 평범하게 풀기 위해 구현을 어떻게 해야 하는지 그 코드를 보도록 하자.

Dynamic Segment Tree Implementation

 

해당 코드는 JusticeHui의 블로그에서 소개된 인덱스 기반 구현이다. 시작 노드가 0부터 시작한다는 점을 주의한다.

 

각 범위에 해당하는 노드를 동적으로 관리하기 위해 구조체 멤버로 l과 r이 있고 쿼리를 처리할 때마다 쓰지 않은 노드를 하나씩 활성화하면서 계산하는 것을 볼 수 있다.

 

자 이제 우리는 코드를 얻었으니 다이나믹 세그먼트 트리를 쓰는 문제를 해결해보자.

 

https://www.acmicpc.net/problem/14577

14577번: 일기예보

 

도입에서 말한 것처럼 무시무시한 범위를 가진 쿼리 문제가 나왔다. 하지만 우리는 다이나믹 세그먼트 트리를 알고 있으므로 그냥 구간 $[0, 10^{18}]$에 대해 쿼리를 잘 수행하기만 하면 된다.

 

라고 하지만 4번 쿼리가 다소 문제다. 4번 쿼리는 kth 쿼리인데 다이나믹 세그먼트 트리로는 어떻게 처리할 수 있는가?

 

그냥 node * 2와 node * 2 + 1 대신 l, r의 원소 개수들로 하여금 k번째 원소를 찾아가면 된다. 정적 세그먼트 트리에서 kth 쿼리를 할 때와 거의 비슷하다.

 

ll kth_query(int node, ll from, ll to, int kth)
{
  if (from == to) return to;
  ll mid = (from + to) >> 1;
  auto lval = tree[tree[node].l].val;
  if (lval >= kth) return kth_query(tree[node].l, from, mid, kth);
  return kth_query(tree[node].r, mid + 1, to, int(kth - lval));
}

 

이 함수를 가지고 4번 쿼리를 처리하면 된다.

 

일기 예보를 해결하는 전체 코드는 다음처럼 짤 수 있다. 노드 개수를 주의하자.

 

#pragma GCC target("sse,sse2,sse3,ssse3,sse4,avx,avx2,fma")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")
 
#include "bits/stdc++.h"
#include "ext/rope"
 
using namespace std;
using namespace __gnu_cxx;
 
using ll = long long;
using pii = pair<int, int>;
 
ll ar[100001];
 
#define SEG_SIZE 3500000
 
struct node {
  int l, r;
  ll val;
  node() { l = r = -1; val = 0; }
};
int tree_size = 1;
 
array<node, SEG_SIZE> tree;
 
void update(int node, ll from, ll to, ll idx, ll val) {
  if (from == to) {
    tree[node].val = val;
    return;
  }
  ll mid = (from + to) >> 1;
  if (idx <= mid) {
    if (tree[node].l < 0) tree[node].l = tree_size++;
    update(tree[node].l, from, mid, idx, val);
  } else {
    if (tree[node].r < 0) tree[node].r = tree_size++;
    update(tree[node].r, mid + 1, to, idx, val);
  }
  int idxl = tree[node].l, idxr = tree[node].r;
  auto vl = idxl != -1 ? tree[idxl].val : 0;
  auto vr = idxr != -1 ? tree[idxr].val : 0;
  tree[node].val = vl + vr;
}
 
ll query(int node, ll from, ll to, ll l, ll r) {
  if (node == -1) return 0;
  if (to < l or from > r) return 0;
  if (l <= from and to <= r) return tree[node].val;
  ll mid = (from + to) >> 1;
  return query(tree[node].l, from, mid, l, r) + query(tree[node].r, mid + 1, to, l, r);
}
 
ll kth_query(int node, ll from, ll to, int kth)
{
  if (from == to) return to;
  ll mid = (from + to) >> 1;
  auto lval = tree[tree[node].l].val;
  if (lval >= kth) return kth_query(tree[node].l, from, mid, kth);
  return kth_query(tree[node].r, mid + 1, to, int(kth - lval));
}
 
int n, q;
 
void solve()
{
  cin >> n >> q;
  for (int i = 1; i <= n; i++)
  {
    cin >> ar[i];
    ll cur = query(0, 0, 1e18, ar[i], ar[i]);
    update(0, 0, 1e18, ar[i], cur + 1);
  }
 
  while (q--)
  {
    int op;
    ll l, r;
    cin >> op >> l;
    if (op == 1)
    {
      cin >> r;
      ll cur = query(0, 0, 1e18, ar[l], ar[l]);
      update(0, 0, 1e18, ar[l], cur - 1);
      ar[l] += r;
      cur = query(0, 0, 1e18, ar[l], ar[l]);
      update(0, 0, 1e18, ar[l], cur + 1);
    }
    else if (op == 2)
    {
      cin >> r;
      ll cur = query(0, 0, 1e18, ar[l], ar[l]);
      update(0, 0, 1e18, ar[l], cur - 1);
      ar[l] -= r;
      cur = query(0, 0, 1e18, ar[l], ar[l]);
      update(0, 0, 1e18, ar[l], cur + 1);
    }
    else if (op == 3)
    {
      cin >> r;
      cout << query(0, 0, 1e18, l, r) << '\n';
    }
    else
    {
      int order = n - l + 1;
      cout << kth_query(0, 0, 1e18, order) << '\n';
    }
  }
}
 
int main()
{
#ifdef LOCAL
  freopen("input.txt", "r", stdin);
//  freopen("output.txt", "w", stdout);
#endif
  cin.tie(nullptr);
  ios::sync_with_stdio(false);
 
  solve();
}

Lazy Propagation with Dynamic Segment Tree

자 그러면 레이지 프로퍼게이션도 가능할까?

 

물론 어렵지 않다. 참고로 이를 Sparse Segment Tree로 부르는 것 같다. 게으른 전파를 하는 다이나믹 세그먼트 트리라고 하면 너무 기니 우리도 Sparse Segment Tree라 부르자.

 

구현도 역시 정적 lazy propagation과 거의 같다. 다음은 구간 덧셈과 구간 합을 구하는 스파스 세그먼트 트리다. 역시 시작 노드는 0이다.

 

 

그러면 스파스 세그먼트 트리로 풀 수 있는 문제를 풀어보자.

 

https://www.acmicpc.net/problem/20212

20212번: 나무는 쿼리를 싫어해~

 

아 역시 이것도 범위가 크다는 것만 빼면 스파스 세그먼트 트리 문제가 된다. 문제에서 대놓고 오프라인 쿼리를 하라고 하므로 쿼리를 잘 정렬한 뒤 처리해주자.

 

나무는 쿼리를 싫어해~를 해결하는 코드는 다음과 같다.

 

#pragma GCC target("sse,sse2,sse3,ssse3,sse4,avx,avx2,fma")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")
 
#include "bits/stdc++.h"
#include "ext/rope"
 
using namespace std;
using namespace __gnu_cxx;
 
using ll = long long;
using pii = pair<int, int>;
 
#define MAXN 200001
 
struct node {
  ll sum, lazy;
  int l, r;
  node() { sum = lazy = 0, l = r = -1; }
};
 
int tree_size = 1;
array<node, MAXN*64> tree;
 
inline int get_mid(int ss, int ee) { return ss + (ee - ss) / 2; }
 
void push(int node, int start, int end) {
  if (node == -1) return;
  auto &lz = tree[node].lazy;
  if (!lz) return;
  tree[node].sum += lz * (end - start + 1);
  if (start != end) {
    int &ptrl = tree[node].l, &ptrr = tree[node].r;
    if (ptrl == -1) ptrl = tree_size++;
    if (ptrr == -1) ptrr = tree_size++;
    tree[ptrl].lazy += lz;
    tree[ptrr].lazy += lz;
  }
  lz = 0;
}
 
void add(int node, int start, int end, int l, int r, ll val) {
  push(node, start, end);
  if (node == -1 or start > r or end < l) return;
  if (l <= start and end <= r) {
    tree[node].lazy += val;
    push(node, start, end);
    return;
  }
  int mid = get_mid(start, end);
  int &ptrl = tree[node].l, &ptrr = tree[node].r;
  if (ptrl == -1) ptrl = tree_size++;
  if (ptrr == -1) ptrr = tree_size++;
  add(ptrl, start, mid, l, r, val);
  add(ptrr, mid + 1, end, l, r, val);
  auto lval = ptrl != -1 ? tree[ptrl].sum : 0;
  auto rval = ptrr != -1 ? tree[ptrr].sum : 0;
  tree[node].sum = lval + rval;
}
 
ll sum(int node, int start, int end, int l, int r) {
  if (node == -1 or start > r or end < l) return 0;
  push(node, start, end);
  if (l <= start and end <= r) return tree[node].sum;
  int mid = get_mid(start, end);
  return sum(tree[node].l, start, mid, l, r) + sum(tree[node].r, mid + 1, end, l, r);
}
 
ll ans[50001];
int q;
int upd_query;
 
struct query
{
  int i, j, k, type, idx, pos;
  bool operator < (query &o) const {
    if (pos != o.pos) return pos < o.pos;
    return type < o.type;
  }
} queries[50000];
 
void solve()
{
  memset(ans, -1, sizeof ans);
  cin >> q;
  for (int i = 0; i < q; i++)
  {
    cin >> queries[i].type >> queries[i].i >> queries[i].j >> queries[i].k;
    queries[i].idx = i;
    if (queries[i].type == 1)
    {
      queries[i].pos = upd_query++;
    }
    else
    {
      queries[i].pos = queries[i].k - 1;
    }
  }
  sort(queries, queries + q);
 
  for (int i = 0; i < q; i++)
  {
    auto [go, to, val, op, idx, priority] = queries[i];
    if (op == 1)
    {
      add(0, 1, 1e9, go, to, val);
    }
    else
    {
      ans[idx] = sum(0, 1, 1e9, go, to);
    }
  }
 
  for (int i = 0; i < q; i++)
  {
    if (ans[i] == -1) continue;
    cout << ans[i] << '\n';
  }
}
 
int main()
{
#ifdef LOCAL
  freopen("input.txt", "r", stdin);
//  freopen("output.txt", "w", stdout);
#endif
  cin.tie(nullptr);
  ios::sync_with_stdio(false);
 
  solve();
}

 

지금까지 다이나믹 세그먼트 트리와 레이지 프로퍼게이션을 적용한 스파스 세그먼트 트리에 대해 알아보았다.

 

이런 세그먼트 트리까지 쓸 정도로 무자비한 문제는 별로 없는 것 같지만 의외로 ICPC도 아니고 xOI에서 나오는 문제다.

 

어쨌든 우리는 든든한 트리를 알게 되었으니 문제를 좀 더 간단하게 풀 수 있을 것이다.

 

다이나믹/스파스 세그먼트 트리로 풀 수 있다고 알려진 문제를 소개하면서 글을 마친다. 풀어보자!

 

그런데 내가 안 해봤으므로 그렇게 풀 수 있는진 확실하지 않다. 그것도 유의하자.

 

https://www.acmicpc.net/problem/11843

https://www.acmicpc.net/problem/11980

https://www.acmicpc.net/problem/11783

https://www.acmicpc.net/problem/5486