[P1] 백준 13925번 수열과 쿼리 13 C++ 세그먼트 트리, 느리게 갱신되는 세그먼트 트리

리뷰

 

쿼리로 부분 수열에 특정 수를 더하기, 곱하기, 대체하기가 나오는 구간합 문제

느리게 갱신되는 세그먼트 트리의 개념을 익혔다고 섣불리 덤볐다가 혼쭐났다.

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

 

문제 풀이

전역 변수

• MAX_N : 노드의 최대 개수를 지정할 정수형 상수 변수
• MOD : MOD 연산을 해줄 정수형 상수 변수, 10억 + 7로 고정이다.
• nodes : 수열의 정보가 저장될 정수형 배열, MAX_N 크기로 초기화 해준다.
• tree : 세그먼트 트리 정보가 저장될 정수형 배열, MAX_N * 4 크기로 초기화 해준다.
• n, m : 수열의 길이를 저장할 변수 n, 쿼리의 개수를 저장할 변수 m
• Lazy : 세그먼트 트리의 구간 더하기(add), 곱하기(mul), 대체하기(set)를 느린 갱신으로 적용하기 위한 구조체, add의 경우 기본값으로 0, mul의 경우 기본값으로 1, set의 경우 기본값으로 -1으로 설정해 줬다.
• lazy : 트리 내에 아직 갱신이 되지 않은 노드를 관리하기 위한 업데이트용 트리, tree와 크기 동일하게 세팅한다.

 

함수

1. build

void build(ll node, ll start, ll end)

 

수열을 세그먼트 트리로 만들어줄 함수

 

2. propagate

void propagate(ll node, ll start, ll end)

 

갱신할 값이 남아있는지 체크 하고, 갱신할 값이 있다면 갱신을 진행해 주는 함수

현재 노드가 리프 노드가 아닐 경우 하위 노드에 갱신 정보를 전파해 준다.

 

3. update

void update(ll node, ll start, ll end, ll L, ll R, ll idx, ll val)

 

구간 업데이트 함수, 입력받은 범위 내 요소들에 대한 업데이트를 진행한다.

idx가 1일 경우 구간에 val만큼의 값을 더해준다.

2일 경우 구간에 val만큼의 값을 곱해주고, 3일 경우 구간을 val로 대체해 준다.

모든 갱신 작업은 propagate에서 진행하니 lazy[node]에 val을 할당해 주고 propagate를 진행해 주면 된다.

 

4. query

ll query(ll node, ll start, ll end, ll L, ll R)

 

구간 내 누적합을 구하고 리턴해 주는 함수

 

문제 풀이

1. n값을 입력 받고, nodes배열에 길이 n의 수열 정보를 입력 받아준다.
2. build함수를 통해 세그먼트 트리를 빌드해 준다.
3. m값을 입력 받고 m만큼의 while문을 실행해 쿼리를 입력 받아준다.
4. a, b, c까지 입력을 받고 만약 a가 3보다 작을 경우 모두 업데이트 작업이므로 d를 추가로 입력 받아준다.
5. update함수에 범위를 매개변수로 넘겨주면서 a와 d를 각각 idx, val로 전달해 준다.
6. a가 4일경우 구간합을 출력할 함수 query를 출력해 준다.

 

참고 사항

• val로 입력받을 수 있는 값의 최대가 10억이므로 long long 타입이어도 곱연산 및 합연산 시 언제든지 터질 가능성이 농후하다.
• 메모리는 어차피 넉넉한 편이니 모든 int를 long long으로 바꾸어 주고 연산이 진행되는 모든 부분에 MOD를 통한 모듈러 연산을 하는 것을 추천한다.

 

 

정답 코드

#include<iostream>
#define ll long long

using namespace std;

const ll MAX_N = 100000;
const ll MOD = 1000000007;
ll nodes[MAX_N];
ll tree[MAX_N * 4];
ll n, m;

struct Lazy {
    ll add = 0;   
    ll mul = 1;   
    ll set = -1;  
};

Lazy lazy[MAX_N * 4];

void build(ll node, ll start, ll end) {
    if (start == end) tree[node] = nodes[start];
    else {
        ll mid = (start + end) / 2;
        build(node * 2, start, mid);
        build(node * 2 + 1, mid + 1, end);
        tree[node] = (tree[node * 2] + tree[node * 2 + 1]) % MOD;
    }
}

void propagate(ll node, ll start, ll end) {
    if (lazy[node].set != -1) {
        tree[node] = (lazy[node].set * (end - start + 1)) % MOD;
        if (start != end) {
            lazy[node * 2].set = lazy[node].set;
            lazy[node * 2 + 1].set = lazy[node].set;
            lazy[node * 2].mul = 1;
            lazy[node * 2 + 1].mul = 1;
            lazy[node * 2].add = 0;
            lazy[node * 2 + 1].add = 0;
        }
        lazy[node].set = -1;
    }

    if (lazy[node].mul != 1) {
        tree[node] = (tree[node] * lazy[node].mul) % MOD;
        if (start != end) {
            lazy[node * 2].mul = (lazy[node * 2].mul * lazy[node].mul) % MOD;
            lazy[node * 2 + 1].mul = (lazy[node * 2 + 1].mul * lazy[node].mul) % MOD;
            lazy[node * 2].add = (lazy[node * 2].add * lazy[node].mul) % MOD;
            lazy[node * 2 + 1].add = (lazy[node * 2 + 1].add * lazy[node].mul) % MOD;
        }
        lazy[node].mul = 1;
    }

    if (lazy[node].add != 0) {
        tree[node] = (tree[node] + lazy[node].add * (end - start + 1)) % MOD;
        if (start != end) {
            lazy[node * 2].add = (lazy[node * 2].add + lazy[node].add) % MOD;
            lazy[node * 2 + 1].add = (lazy[node * 2 + 1].add + lazy[node].add) % MOD;
        }
        lazy[node].add = 0;
    }
}

void update(ll node, ll start, ll end, ll L, ll R, ll idx, ll val) {
    propagate(node, start, end);
    if (R < start || end < L) return;
    if (L <= start && end <= R) {
        if (idx == 1) lazy[node].add = val;
        else if (idx == 2) lazy[node].mul = val;
        else lazy[node].set = val;
        propagate(node, start, end);
        return;
    }
    ll mid = (start + end) / 2;
    update(node * 2, start, mid, L, R, idx, val);
    update(node * 2 + 1, mid + 1, end, L, R, idx, val);
    tree[node] = (tree[node * 2] + tree[node * 2 + 1]) % MOD;
}

ll query(ll node, ll start, ll end, ll L, ll R) {
    propagate(node, start, end);
    if (R < start || end < L) return 0;
    if (L <= start && end <= R) return tree[node] % MOD;
    ll mid = (start + end) / 2;
    ll left = query(node * 2, start, mid, L, R);
    ll right = query(node * 2 + 1, mid + 1, end, L, R);
    return (left + right) % MOD;
}

int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);

    cin >> n;
    for (ll i = 0; i < n; i++) cin >> nodes[i];
    build(1, 0, n - 1);
    cin >> m;
    while (m--) {
        ll a, b, c; cin >> a >> b >> c;
        if (a <= 3) {
            ll d; cin >> d;
            update(1, 0, n - 1, b - 1, c - 1, a, d);
        }
        else cout << query(1, 0, n - 1, b - 1, c - 1) % MOD << "\n";
    }
}