【树状数组】POJ 3321 Apple Tree

/**
 *  @author        johnsondu
 *  @time          2015.8.25 20:04
 *  @problem       POJ 3321 Apple Tree
 *  @type          Binary Index Tree
 *  @description   从根节点开始，dfs遍历树，先访问的节点
 *                 记为beg, 从当前结点遍历访问到的最后的
 *                 一个节点，记为end。然后按照树状数组的
 *                 方法进行求解。
 *  @url           http://poj.org/problem?id=3321
 */

#include <iostream>
#include <cstdio>
#include <cmath>
#include <cstdlib>
#include <algorithm>
#include <cstring>
#include <string>
#include <stack>
#include <queue>
#include <map>
#include <exception>
#include <vector>
#include <stdexcept>

using namespace std;

const int N = 100005;

struct TNode
{
    int to, next;
}node[2 * N];

int m, n, cnt;
bool vis[N], apple[N];
int beg[N], en[N], p[N], tot[N];

void init()
{
    cnt = 0;
    memset(p, -1, sizeof(p));
    memset(tot, 0, sizeof(tot));
    memset(vis, false, sizeof(vis));
    memset(apple, true, sizeof(apple));
}

void addedge(int u, int v)
{
    node[cnt].to = v;
    node[cnt].next = p[u];
    p[u] = cnt ++;
}

void dfs(int u)
{
    beg[u] = ++ cnt;
    vis[u] = true;
    for(int i = p[u]; i != -1; i = node[i].next) {
        if(!vis[node[i].to]) {
            dfs(node[i].to);
        }
    }
    en[u] = cnt;
}

int lowbit(int x)
{
    return  x & (-x);
}

void update(int x, int val)
{
    for(; x <= n; x += lowbit(x))
        tot[x] += val;
}

int getSum(int x)
{
    int sum = 0;
    for(; x > 0; x -= lowbit(x))
        sum += tot[x];
    return sum;
}

int main()
{
    int u, v;
    char str[10];
    init();
    scanf("%d", &n);
    for(int i = 1; i < n; i ++) {
        scanf("%d%d", &u, &v);
        addedge(u, v);
        addedge(v, u);
    }

    // rearange the order
    cnt = 0;
    dfs(1);

    // firstly, all branches have apples
    for(int i = 1; i <= n; i ++){
        update(i, 1);
    }

    scanf("%d", &m);
    while(m --) {
        scanf("%s%d", str, &u);
        if(str[0] == ‘C‘) {
            if(!apple[u]) {
                update(beg[u], 1);
                apple[u] = true;
            }
            else {
                update(beg[u], -1);
                apple[u] = false;
            }
        }
        else {
            printf("%d\n", getSum(en[u]) - getSum(beg[u]-1));
        }
    }


    return 0;
}