세그먼트 트리 with Lazy propagation

저번 글에서 점 갱신과 구간 연산을 \(O(\log n)\)에 처리할 수 있는 자료구조인 세그먼트 트리에 대해 알아봤습니다.

이번엔 구간 갱신까지 \(O(\log n)\)에 처리하게 해주는 Lazy propagation에 대해 알아봅시다.

 

Lazy propagation을 하기 앞서, 다음 문제를 풀어봅시다.

 

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

16975번: 수열과 쿼리 21

이 문제처럼 구간 갱신이 필요하지만, 구간에 대한 연산결과가 아닌 한 점의 값만 필요한 경우에는, Lazy propagation없이 일반 세그먼트 트리만을 이용해 쿼리를 처리할 수 있습니다.

 

다음과 같은 수열이 있다면, 수열의 값 자체를 저장하는 것이 아니라 이전 값과 비교했을 때의 변화량을 저장합시다.

변화량을 저장한 수열로 세그먼트 트리를 만듭시다.

그러면 \(x\)번째 원소의 값을 구하는 연산은 \([1,x]\)구간의 합으로 구할 수 있게 됩니다.

 

그렇다면 구간 갱신은 어떻게 처리할 수 있을까요?

위 수열 A에서 3번째 원소에서 5번째 원소까지 구간에 2를 더한다고 생각해봅시다.

변화량이 바뀌는 것은 2곳 뿐입니다. 3번째 변화량에 2가 더해졌고, 6번째 변화량에 -2가 더해졌습니다.

 

일반적으로 구간 \([i,j]\)에 값 \(x\)를 더한다면, 변화량 수열에서 \(i\)번째 값에 \(x\)를 더하고, \(j+1\)번째 값에 \(x\)를 뺌으로써 처리할 수 있습니다.

 

따라서 이 문제도 점 갱신과 구간 연산을 지원하는 일반 세그먼트 트리로 풀 수 있게 됩니다.

 

#include <bits/stdc++.h>
using namespace std;
 
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
 
ll gcd(ll a, ll b) { for (; b; a %= b, swap(a, b)); return a; }
 
const int N = 100001;
 
int n;
ll segTree[N * 4];
 
void update(int ptr, int l, int r, int i, ll val)
{
    if (i < l || r < i) return;
    if (l == r)
    {
        segTree[ptr] += val;
        return;
    }
 
    update(ptr * 2, l, (l + r) / 2, i, val);
    update(ptr * 2 + 1, (l + r) / 2 + 1, r, i, val);
 
    segTree[ptr] = segTree[ptr * 2] + segTree[ptr * 2 + 1];
}
 
ll getVal(int ptr, int l, int r, int i, int j)
{
    if (j < l || r < i) return 0;
    if (i <= l && r <= j) return segTree[ptr];
 
    return getVal(ptr * 2, l, (l + r) / 2, i, j)
        + getVal(ptr * 2 + 1, (l + r) / 2 + 1, r, i, j);
}
 
int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0), cout.tie(0);
 
    cin >> n;
    ll last = 0;
    for (int i = 1; i <= n; i++)
    {
        ll a; cin >> a;
        update(1, 1, n, i, a - last);
        last = a;
    }
 
    int m; cin >> m;
    while (m--)
    {
        int q; cin >> q;
        if (q == 1)
        {
            ll i, j, k; cin >> i >> j >> k;
            update(1, 1, n, i, k);
            update(1, 1, n, j + 1, -k);
        }
        else
        {
            ll x; cin >> x;
            cout << getVal(1, 1, n, 1, x) << '\n';
        }
    }
}

펜윅트리를 이용해서도 쉽게 구현할 수 있습니다.

 

#include <bits/stdc++.h>
using namespace std;
 
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
 
ll gcd(ll a, ll b) { for (; b; a %= b, swap(a, b)); return a; }
 
const int N = 100001;
 
int n;
ll segTree[N];
 
void update(int i, ll val)
{
    while (i <= n)
    {
        segTree[i] += val;
        i += i & -i;
    }
}
 
ll getVal(int i)
{
    ll res = 0;
    while (i)
    {
        res += segTree[i];
        i -= i & -i;
    }
    return res;
}
 
int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0), cout.tie(0);
 
    cin >> n;
    ll last = 0;
    for (int i = 1; i <= n; i++)
    {
        ll a; cin >> a;
        update(i, a - last);
        last = a;
    }
 
    int m; cin >> m;
    while (m--)
    {
        int q; cin >> q;
        if (q == 1)
        {
            ll i, j, k; cin >> i >> j >> k;
            update(i, k);
            update(j + 1, -k);
        }
        else
        {
            ll x; cin >> x;
            cout << getVal(x) << '\n';
        }
    }
}

 

이 방법은 일반적으로 역원이 존재하는 연산에 대해서만 사용할 수 있음을 알아둡시다.

---

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

10999번: 구간 합 구하기 2

이제 가장 일반적인 경우에 대해 살펴봅시다.

구간에 갱신과, 구간에 연산 결과를 모두 빠르게 처리해야 합니다.

 

구간에 대한 연산은 일반 세그먼트 트리와 같으니, 구간 갱신을 어떻게 빠르게 할 수 있는지 알아봅시다.

 

다음과 같은 세그먼트 트리가 있다고 합시다.

여기에서 구간 \([2,6]\)에 3을 더해주는 연산을 한다고 했을 때, 세그먼트 트리에서 갱신되어야 하는 원소는 다음과 같습니다.

 

 

갱신해야 할 원소가 너무 많습니다.

최악의 경우에는 세그먼트 트리 전체를 갱신해야 하므로, 모든 원소를 그때그때 갱신해주는것은 좋지 않은 방법입니다.

 

우리가 구간 연산결과를 계산하는 방법으로 구간 \([2,6]\)의 합을 처리하기 위해 더해줘야 하는 원소는 다음과 같습니다.

이 구간을 갱신해 나갈 때, 표시한 원소까지만 확실히 갱신 해준다면, 전체적인 결과는 달라지지 않으면서도 최소한의 갱신 수를 이용해 갱신할 수 있음을 알 수 있습니다.

 

따라서 다음과 같이 갱신할 수 있습니다.

회색 : 아직 갱신되지 않은 원소

 

일단 이러한 방법으로 최대 \(O(\log n)\)번의 갱신만으로도 구간을 갱신할 수 있음을 알 수 있습니다.

 

아직 갱신되지 않은 자식 원소들은 나중에 여기에 3을 더해줘야 한다는 사실을 어딘가에 저장해 놓은 다음,

나중에 이 원소의 값을 따로 계산해야 할 때가 되면 그 때 한번에 처리해 주면 됩니다.

 

따라서 Lazy 트리를 하나 만든다음, 이 원소(와 그 자식 노드)에 3을 더해야 한다는 의미로 3을 저장해 놓읍시다.

 

이 Lazy 트리에 저장된 값은 해당하는 원소의 값을 직접 사용해야 할 때, 원래 세그먼트 트리에 처리해 주면 됩니다.

 

이로써 구간 갱신과 구간 연산을 모두 쿼리당 \(O(\log n)\)에 처리할 수 있게 됩니다.

 

#include <bits/stdc++.h>
using namespace std;
 
typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
 
ll gcd(ll a, ll b) { for (; b; a %= b, swap(a, b)); return a; }
 
const int N = 1000001;
 
int n, m, k;
ll segTree[N * 4];
ll lazy[N * 4]; // Lazy 트리
 
void setLazy(int ptr, int l, int r)
{
    // 현재 원소에 저장된 Lazy 값을 지금 처리해준다.
    ll val = lazy[ptr];
    lazy[ptr] = 0;
 
    segTree[ptr] += (r - l + 1) * val;
 
    if (l != r)
    {
        // 리프노드가 아니라면, 이 값을 또 자식 노드의 Lazy배열로 밀어준다.
        lazy[ptr * 2] += val;
        lazy[ptr * 2 + 1] += val;
    }
}
 
void update(int ptr, int l, int r, int i, int j, ll val)
{
    if (lazy[ptr]) setLazy(ptr, l, r);
    // lazy 트리에 뭔가 저장되어 있다면 지금 처리한다.
 
    if (j < l || r < i) return;
    if (i <= l && r <= j)
    {
        segTree[ptr] += (r - l + 1) * val;
 
        if (l != r)
        {
            // 리프노드가 아니라면, 이 값을 자식 노드의 Lazy배열로 밀어준다.
            lazy[ptr * 2] += val;
            lazy[ptr * 2 + 1] += val;
        }
        return;
    }
 
    update(ptr * 2, l, (l + r) / 2, i, j, val);
    update(ptr * 2 + 1, (l + r) / 2 + 1, r, i, j, val);
 
    segTree[ptr] = segTree[ptr * 2] + segTree[ptr * 2 + 1];
}
 
ll getVal(int ptr, int l, int r, int i, int j)
{
    if (lazy[ptr]) setLazy(ptr, l, r);
    // lazy 트리에 뭔가 저장되어 있다면 지금 처리한다.
 
    if (j < l || r < i) return 0;
    if (i <= l && r <= j) return segTree[ptr];
 
    return getVal(ptr * 2, l, (l + r) / 2, i, j)
        + getVal(ptr * 2 + 1, (l + r) / 2 + 1, r, i, j);
}
 
int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0), cout.tie(0);
 
    cin >> n >> m >> k;
    for (int i = 1; i <= n; i++)
    {
        ll x; cin >> x;
        update(1, 1, n, i, i, x);
    }
 
    for (int i = 0; i < m + k; i++)
    {
        int q; cin >> q;
        if (q == 1)
        {
            ll b, c, d; cin >> b >> c >> d;
            update(1, 1, n, b, c, d);
        }
        else
        {
            ll b, c; cin >> b >> c;
            cout << getVal(1, 1, n, b, c) << '\n';
        }
    }
}

---

제 그룹의 문제집에서 연습 문제들을 관리하고 있습니다.
문제집의 문제들을 보고 싶으시다면, 가입 신청을 해 주세요.

 

www.acmicpc.net/group/7712

ANZ1217