백준 18227번 : 성대나라의 물탱크 (c++)

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

18227번: 성대나라의 물탱크

문제

성대나라에는 각 도시별로 가뭄을 대비하기 위한 물탱크가 하나씩 존재한다. 이 물탱크들은 모두 연결되어있으며, 루트(성대나라의 수도)가 있는 트리의 형태를 가진다.
지금 성대나라는 물탱크의 물을 사용하여 가뭄을 버텨냈으나, 그 영향으로 모든 물탱크에 물이 비어버리고 말았다.

성대나라의 물관리 시스템은 다소 특수해서, 물은 항상 다음과 같은 방식으로 채워진다:

A번 도시에 물을 채우기로 했다면, 수도에서부터 A번 도시까지 잇는 직선 경로에
수도부터 차례대로 1L, 2L, ⋯이 채워져서 A번 도시에는 (수도부터 A번 도시까지의 도시 수) L 만큼 추가된다.

예를 들어, 아래 그림과 같이 물탱크가 연결되어 있을 때, "4번 도시에 물을 채운다"라고 하면, 1번 도시에 1L, 4번 도시에 2L의 물이 추가된다. 만약 "5번 도시에 물을 채운다"라고 하면 1번 도시에 1L, 2번 도시에 2L, 5번 도시에 3L의 물이 추가된다.

성대나라의 물탱크 관리 담당인 균관이는 어느 도시에 몇 리터의 물이 저장되어있는지 자신이 궁금해질 때마다 알기를 원한다. 균관이를 도와주는 프로그램을 만들어보자.

입력
첫째 줄에 성대나라의 도시의 수 N (1 ≤ N ≤ 200,000)과 수도의 번호 C (1 ≤ C ≤ N)가 공백으로 구분되어 주어진다.

둘째 줄부터 N-1개의 줄에 연결되어있는 두 도시의 번호 쌍 x, y가 공백으로 구분되어 주어진다(1 ≤ x, y ≤ N, x ≠ y). 물탱크의 연결 형태는 트리 구조임이 보장된다. N+1번째 줄에 질의의 수 Q(1 ≤ Q ≤ 200,000)가 주어진다. N+2번째 줄부터 Q개의 줄에 질의가 들어온다. 질의는 다음과 같이 두 종류 중 하나로 주어진다:

• 1 A : A도시에 물을 채운다.

• 2 A : 현재 A도시에 얼마만큼의 물이 채워져 있는지 출력하라.

두 경우 모두 1 ≤ A ≤ N  을 만족한다.

출력
2로 시작하는 질의가 올 때 마다 그 결과를 한 줄에 하나씩 출력한다.

예제 입력 1 
7 1
1 2
1 3
1 4
2 5
4 6
4 7
8
1 6
2 1
2 4
2 6
1 1
2 1
2 4
2 6
예제 출력 1 
1
2
3
2
2
3

#include <iostream>
#include <algorithm>
#include <cmath>
#include <vector>

using namespace std;
typedef long long ll;
const int MAX = 200001;
int N, C, Q;
vector <ll> num2;
vector<int> node[MAX];
int d[MAX];
struct tree {
    int start;
    int end;
    int index;
};
tree trees[MAX];

struct IndexedTree {
    vector <ll> tree;
    vector <ll> num;
    int leafSize = 0, depth = 0;

    IndexedTree(vector <ll> num1) {
        num = num1;
        while (pow(2, depth) < num1.size() - 1) {
            depth++;
        }
        leafSize = pow(2, depth);
        while (num.size() <= leafSize) {
            num.push_back(0);
        }
        tree.resize(2 * leafSize);
    }

    ll query(int node, int left, int right, int qLeft, int qRight) {
        if (qRight < left || right < qLeft) {
            return 0;
        }
        else if (qLeft <= left && right <= qRight) {
            return tree[node];
        }
        else {
            ll mid = (left + right) / 2;
            return query(node * 2, left, mid, qLeft, qRight) + query(node * 2 + 1, mid + 1, right, qLeft, qRight);
        }
    }

    void update(int node, int left, int right, int index) {
        if (index < left || right < index) {
            return;
        }
        else {
            tree[node] += 1;
            if (left != right) {
                ll mid = (left + right) / 2;
                update(2 * node, left, mid, index);
                update(2 * node + 1, mid + 1, right, index);
            }
        }
    }
};


int cnt = 1;
void dfs(int n){
    trees[n].start = cnt;
    trees[n].index = cnt;
    cnt++;
    for (int i = 0; i < node[n].size(); i++) {
        int ny = node[n][i];
        if (!d[ny]) {
            d[ny] = d[n] + 1;
            dfs(ny);
        }
    }
    trees[n].end = cnt-1;
}


int main() {

    ios_base::sync_with_stdio(0);
    cin.tie(0); cout.tie(0);
    freopen("input.txt", "r", stdin);
    
    num2.resize(200001);
    IndexedTree h(num2);

    cin >> N >> C;
    for (int i = 0; i < N-1; i++)
    {
        int a, b;
        cin >> a >> b;
        node[a].push_back(b);
        node[b].push_back(a);
    }
    d[C] = 1;
    dfs(C);

    cin >> Q;

    for (int i = 0; i < Q; i++)
    {
        int a, b;
        cin >> a >> b;
        if (a == 1) {
            h.update(1, 1, h.leafSize, trees[b].index);
        }
        else if (a == 2) {
           cout << h.query(1, 1, h.leafSize, trees[b].start, trees[b].end)*d[b]<<"\n";
        }
    }
    return 0;
}