[c++] 백준 내리갈굼 시리즈 with 세그먼트 트리 (190906)

내리 갈굼 https://www.acmicpc.net/problem/14267

내리 갈굼 2 https://www.acmicpc.net/problem/14268

내리 갈굼 3 https://www.acmicpc.net/problem/14287

내리 갈굼 4 https://www.acmicpc.net/problem/14288

 

세그먼트 트리 관련 문제이다. 1, 2는 세그먼트 트리를 곧바로 사용해도 되지만, 야자타임이 발생하는 3, 4번은 약간 응용 같은 느낌이다. 3번은 개별 update, 범위 find를 이용하면 된다. 

4번은 개별 tree와 범위 tree 두 개를 동시에 쓰면 된다. 이 때 트리를 별도로 만들어야 한다. 왜냐하면 한 트리에서 야자타임했다가 안했다가 하면 구간합과 개별합이 엉망으로 합쳐지므로..

4번의 소스코드만 첨부하면 다음과 같다.

#include <bits/stdc++.h>
using namespace std;
#define INF 987654321
typedef long long ll;
typedef pair<int, int> ii;
typedef tuple<int, int, int> iii;

#define pb push_back
#define MAX_N 100100
int in[MAX_N];
int out[MAX_N];
int trace_num = 0;
vector<int> edge[MAX_N];


int N, Q;

void dfs(int start){
    in[start] = ++trace_num;
    for(auto next : edge[start]){
        if(in[next]) continue; // 이미 방문한 노드
        dfs(next);
    }
    out[start] = trace_num;
}

long long init(vector<long long> &a, vector<long long> &tree, int node, int start, int end) {
    if (start == end) {
        return tree[node] = a[start];
    } else {
        return tree[node] = init(a, tree, node*2, start, (start+end)/2) + init(a, tree, node*2+1, (start+end)/2+1, end);
    }
}
void update_lazy(vector<long long> &tree, vector<long long> &lazy, int node, int start, int end) {
    if (lazy[node] != 0) {
        tree[node] += (end-start+1)*lazy[node];
        // leaf가 아니면 아래로 내려버린다.
        if (start != end) {
            lazy[node*2] += lazy[node];
            lazy[node*2+1] += lazy[node];
        }
        lazy[node] = 0;
    }
}
void update_range(vector<long long> &tree, vector<long long> &lazy, int node, int start, int end, int left, int right, long long diff) {
    update_lazy(tree, lazy, node, start, end);
    if (left > end || right < start) {
        return;
    }
    if (left <= start && end <= right) {
        tree[node] += (end-start+1)*diff;
        if (start != end) {
            lazy[node*2] += diff;
            lazy[node*2+1] += diff;
        }
        return;
    }
    update_range(tree, lazy, node*2, start, (start+end)/2, left, right, diff);
    update_range(tree, lazy, node*2+1, (start+end)/2+1, end, left, right, diff);
    tree[node] = tree[node*2] + tree[node*2+1];
}
long long sum(vector<long long> &tree, vector<long long> &lazy, int node, int start, int end, int left, int right) {
    update_lazy(tree, lazy, node, start, end);
    if (left > end || right < start) {
        return 0;
    }
    if (left <= start && end <= right) {
        return tree[node];
    }
    return sum(tree, lazy, node*2, start, (start+end)/2, left, right) + sum(tree, lazy, node*2+1, (start+end)/2+1, end, left, right);
}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);
    int N, M;
    cin >> N >> M;
    int boss = 0;
    // make a graph
    for(int i=1; i<=N;i++){
        int sv; cin >> sv;
        if (sv == -1) boss = i;
        edge[sv].pb(i); // superviser부터 directed graph
    }
    dfs(boss);

    int n = 3*N; // 모든 노드가 2번씩 나오니까(dfs)
    vector<ll> a(n);
    int h = (int)ceil(log2(n));
    int tree_size = (1 << (h+2)) - 1;
    vector<ll> tree(tree_size);
    vector<ll> lazy(tree_size);
    
    vector<ll> yaza_tree(tree_size);
    vector<ll> yaza_lazy(tree_size);

    init(a, tree, 1, 0, n-1); // 모두 0으로 초기화(맨 처음에는 아무도 안갈궈진 상태)
    int p, i, w;
    bool yaza_time = false;
    while(M--){
        cin >> p;
        if (p == 1){ // i번째 직원이 갈굼당함
            cin >> i >> w;
            int start = in[i]; 
            int end = out[i];
            if (yaza_time){ // 야자타임 트리
                update_range(yaza_tree, yaza_lazy, 1, 0, n-1, start-1, start-1, w); 
            }
            else { // 일반 트리
                update_range(tree, lazy, 1, 0, n-1, start-1, end-1, w);
            }
        }
        else if (p == 2) { // i번째 직원이 갈굼을 당한 정도를 출력 : 부하의 값을 다 본다
            cin >> i;
            int start = in[i]; 
            int end = out[i];
            cout << sum(yaza_tree, yaza_lazy, 1, 0, n-1, start-1, end-1) + sum(tree, lazy, 1, 0, n-1, start-1, start-1) << '\n';
        }
        else { // 야자타임
            if(yaza_time){
                yaza_time = false;
            }
            else yaza_time = true;
        }
    }
    return 0;
}