[C++] 백준 10999 : 구간 합 구하기 2

1. 문제 출처

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

10999번: 구간 합 구하기 2

 

 

 

 

 

 

2. 문제 접근

 

기존에 알고 있었던 세그먼트 트리이지만, 업데이트해야 하는 것이 하나의 값이 아니라 구간에 대해서 업데이트를 진행해야 했다. 해당 구간에 대해서 일일이 업데이를 진행하는 것보다 lazy propataion방법을 사용하면 더 효율적으로 풀이가 가능했다. lazy propagation에 대한 정보가 없었던 나는 다른 블로그들을 참고하여서 문제 풀이를 진행했다.

 

 

 

 

lazy propagation기법은 꽤 단순한 방법이다. 업데이트하고자 하는 구간에 lazy라는 값을 사용한다. 특정 구간을 포함하는 노드가 있다고 하면 그 하위 노드까지 전부 전파를 하는 것이 아니라 구간을 포함하는 노드에만 lazy를 부여하는 것이다. 그렇다면 하위 노드의 값을 구하고 싶을 때, 전파가 안되어 있으면 올바르지 않은 값을 구하는 것이 아닐까?라는 걱정을 했었다. 하지만 이는 query함수를 실행할 때, 어차피 최상위 노드 1번 노드부터 내려오기 때문에 하위 노드를 가기 위해서는 구하고 싶은 구간을 포함하는 중간 노드를 거쳐야만 한다. 그때, 하위 노드로 전파를 시켜주는 것이다.

 

 

 

 

일반 세그먼트 트리의 update 함수는 업데이트하고자 하는 값의 index와 diff값이 인자로 넘어갔다면, 이번 lazy propagation에서는 index가 아닌 그 구간이 인자로 주어진다. query함수와 비슷하게 내가 변경하고 싶은 구간이 이번 노드의 구간이 아니면 update하지 않도록 주건을 부여하면 완료된다.

 

 

 

 

3. 코드

#include <bits/stdc++.h>
#define ll long long
#define pii pair<int, int>

using namespace std;

int n, m, k; // 수의 개수, 변경 횟수, 구간합 쿼리 횟수
vector<ll> elem;
vector<pair<ll, ll>> tree; // 구간합, lazy값


ll initTree(int node, int left, int right) {
    if (left == right) return tree[node].first = elem[left];
    else {
        int mid = (left + right) / 2;
        return tree[node].first = initTree(node * 2, left, mid) + initTree(node * 2 + 1, mid + 1, right);
    }
}


// 구하고 싶은 구간: [start, end]
// 현재 노드가 담고 있는 구간: [left, right]
ll query(int node, int left, int right, int start, int end) {
    // 현재 노드에 lazy값이 남아서 하위 트리에도 전파를 해주어야 한다.
    if (tree[node].second != 0) {
        // 현재 노드의 값을 업데이트 한다.
        tree[node].first += (right - left + 1) * tree[node].second;
        if (left != right) {
            // 자식에게도 값을 업데이트 해준다.
            tree[node * 2].second += tree[node].second;
            tree[node * 2 + 1].second += tree[node].second;
        }
        tree[node].second = 0;            
    }

    if (right < start || end < left) return 0;

    if (start <= left && right <= end) return tree[node].first;

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


void update(int node, int left, int right, int start, int end, ll diff) {
    if (tree[node].second != 0) {
        tree[node].first += (right - left + 1) * tree[node].second;
        if (left != right) {
            tree[node * 2].second += tree[node].second;
            tree[node * 2 + 1].second += tree[node].second;
        }
        tree[node].second = 0;
    } 

    if (end < left || right < start) return;

    if (start <= left && right <= end) {
        tree[node].first += (right - left + 1) * diff;
        if (left != right) {
            tree[node * 2].second += diff;
            tree[node * 2 + 1].second += diff;
        }
        return;
    }

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

    tree[node].first = tree[node * 2].first + tree[node * 2 + 1].first;
}
 
 
int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    cin >> n >> m >> k;
    elem.resize(n);
    for (auto &e : elem) cin >> e;

    int treeHeight = ceil(log2(n));
    int treeSize = 1 << (treeHeight + 1);
    tree.resize(treeSize, {0, 0});
    
    initTree(1, 0, n - 1);

    for (int queryCnt = 0; queryCnt < m + k; ++queryCnt) {
        int opt, start, end; cin >> opt >> start >> end;
        if (opt == 1) { // 구간을 업데이트 하는 경우
            ll diff; cin >> diff;
            update(1, 0, n - 1, start - 1, end - 1, diff);
        } else if(opt == 2) { // 구간합을 질의 하는 경우
            cout << query(1, 0, n - 1, start - 1, end - 1) << "\n";
        }
    }

    return 0;
}

 

 

 

 

 

4. 참조

https://bowbowbow.tistory.com/4

Segment Tree and Lazy Propagation