[C++] 수열과 쿼리 13 (백준 13925번)

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

 

쿼리는 원소들의 합 구하는 쿼리 한 종류이고,

갱신은 원소에 값 더하기, 원소에 값 곱하기 두 종류가 있다. (세번째 쿼리는 그냥 0을 곱한 다음에 v를 더해준다고 생각할 수 있기 때문에 따로 생각해줄 필요는 없다.)

 

더하기와 곱하기가 함께 있는 경우, 어떻게 세그먼트 트리를 관리해야 할까?

 

우선 ax + b 형태로 lazy 배열을 관리해주어야 한다. 나는 lazya와 lazyb로 관리해주었다.

 

일단 update 함수는 쉽다.

더하는 갱신이라면 lazyb에만 값을 더해준다. 왜냐하면 (ax + b) + b' = ax + (b + b') 이기 때문이다.

곱하는 갱신이라면 lazya와 lazyb 각각에 값을 곱해준다. 왜냐하면 (ax + b) * a' = aa'x + ba' 이기 때문이다.

 

push (propagate) 함수는 살짝 복잡하다.

부모의 lazya 값과 lazyb 값을 각각 ap, bp로, 자식의 lazya 값과 lazyb 값을 각각 ac, bc라고 하면 아래와 같이 생각할 수 있다.

v → ac * v + bc → ap * (ac * v + bc) + bp = (ap * ac) * v + (ap * bc + bp)

즉, 아래와 같이 갱신하면 된다.

ac ← ap * ac

bp ← ap * bc + bp

 

#include <iostream>
#include <vector>
using namespace std;

long long MOD = 1e9 + 7;
int n, m;

struct SegTree {
    vector<long long> tree;
    vector<long long> lazya;
    vector<long long> lazyb;

    void push(int idx, int l, int r) {
        if (lazya[idx] == 1 && lazyb[idx] == 0) return;
        tree[idx] = (tree[idx] * lazya[idx] + (r - l + 1) * lazyb[idx]) % MOD;
        if (l < r) {
            lazya[idx * 2] = (lazya[idx * 2] * lazya[idx]) % MOD;
            lazyb[idx * 2] = (lazya[idx] * lazyb[idx * 2] + lazyb[idx]) % MOD;
            lazya[idx * 2 + 1] = (lazya[idx * 2 + 1] * lazya[idx]) % MOD;
            lazyb[idx * 2 + 1] = (lazya[idx] * lazyb[idx * 2 + 1] + lazyb[idx]) % MOD;
        }
        lazya[idx] = 1;
        lazyb[idx] = 0;
    }

    long long update1(int idx, int l, int r, int wl, int wr, long long val) {
        push(idx, l, r);
        if (wr < l || wl > r) return tree[idx];
        if (wl <= l && wr >= r) {
            lazyb[idx] = (lazyb[idx] + val) % MOD;
            push(idx, l, r);
            return tree[idx];
        }
        int m = (l + r) / 2;
        return tree[idx] = (update1(idx * 2, l, m, wl, wr, val) + update1(idx * 2 + 1, m + 1, r, wl, wr, val)) % MOD;
    }

    long long update2(int idx, int l, int r, int wl, int wr, long long val) {
        push(idx, l, r);
        if (wr < l || wl > r) return tree[idx];
        if (wl <= l && wr >= r) {
            lazya[idx] = (lazya[idx] * val) % MOD;
            lazyb[idx] = (lazyb[idx] * val) & MOD;
            push(idx, l, r);
            return tree[idx];
        }
        int m = (l + r) / 2;
        return tree[idx] = (update2(idx * 2, l, m, wl, wr, val) + update2(idx * 2 + 1, m + 1, r, wl, wr, val)) % MOD;
    }
    
    long long query(int idx, int l, int r, int wl, int wr) {
        push(idx, l, r);
        if (wr < l || wl > r) return 0;
        if (wl <= l && wr >= r) return tree[idx];
        int m = (l + r) / 2;
        return (query(idx * 2, l, m, wl, wr) + query(idx * 2 + 1, m + 1, r, wl, wr)) % MOD;
    }
    
    long long update1(int wl, int wr, long long val) {
        return update1(1, 0, n, wl, wr, val);
    }

    long long update2(int wl, int wr, long long val) {
        return update2(1, 0, n, wl, wr, val);
    }
    
    long long query(int wl, int wr) {
        return query(1, 0, n, wl, wr);
    }
};

void solve() {
    SegTree seg;

    cin >> n;
    seg.tree.assign(4 * n + 1, 0);
    seg.lazya.assign(4 * n + 1, 1);
    seg.lazyb.assign(4 * n + 1, 0);

    for (int i = 1; i < n + 1; i++) {
        long long x;
        cin >> x;
        seg.update1(i, i, x);
    }

    cin >> m;
    while (m--) {
        int q, x, y;
        long long v;
        cin >> q;
        if (q == 1) {
            cin >> x >> y >> v;
            seg.update1(x, y, v);
        } else if (q == 2) {
            cin >> x >> y >> v;
            seg.update2(x, y, v);
        } else if (q == 3) {
            cin >> x >> y >> v;
            seg.update2(x, y, 0);
            seg.update1(x, y, v);
        } else if (q == 4) {
            cin >> x >> y;
            cout << seg.query(x, y) << '\n';
        }
    }
}

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

    int T = 1;
    for (int i = 0; i < T; i++) {
        solve();
    }
}

(AC)

 

---

 

내 첫 다이아 문제!

사실 원래는 레이지 날먹 하고 싶어서 (...) 레이지 문제 중에서 푼 사람 수 많은 거 티어 안 보고 암거나 뽑은거라 다이아인 줄도 몰랐다.