POJ 3321 Apple Tree 树状数组题解

原题链接

Apple Tree

• 题意

• 一棵多叉树每个结点有一个编号和一个值，在已知树的结构的情况下，进行两种操作。第一种是修改编号对应的结点的值，第二种是查询某个子树的值的和val。结点数量N<=100000，操作次数M<=100000。

• 朴素的想法是直接模拟，修改的时间复杂度是O(1)，然而查询的时间复杂度是O(N)，进行M次查询可能出现最坏的情况是O(MN)，会超时。

• 考虑到，查询的时间复杂度达到O(N)的主要原因在于，每次查询都遍历了一棵子树，要想办法降低这个时间复杂度。于是思考，能否用一条线段来代表一棵子树，使得对线段的求和就是子树的和val呢？如果我们用线段的左端点表示这棵树，那么该点所有的后代节点都应该在这条线段内。于是可以深度优先搜索，并按访问次序编号，这样一个点的所有后代节点的编号都紧跟在该节点之后，形成了一条线段。可以在深搜时用一个数组记录下这条线段的长度。

• 接下来只要使用树状数组来辅助求和，就可以把时间复杂度降低到O(MlogN)。

• 提交后发现Wrong Answer，原因是修改树状数组的函数写错了，k是向后增加的，而不是像求和一样向前。

• 另外，我第一次做这题的时候用的是线段树，当时的解题报告如下。

• 修改一个结点就要更改它的所有祖先的值val，这种结构类似于线段树。能否将每个结点映射到一条线段呢？这样的映射应该满足如下条件：
  1.该线段包含于其父结点的线段
  2.该结点的所有孩子包含于线段内
  3.该线段中应该有这个结点的值的信息
  4.该线段中应该有val信息

• 于是有一个巧妙的思路，深搜这棵树，按访问顺序重新编号线段的左值，按孩子中最大的右值编号右值，若无孩子，则令右值等于左值。

• 这样，当我们进行修改操作时，只需要修改深搜编号位置处的值，然后递归地修改每一层线段树。当我们进行查询操作时，只需要利用线段树查询这个区间的和。

• 要注意映射关系：每个结点对应于一个编号（线段左值），每个结点也同时对应于一个查询区间（线段）

• 将任务分解如下：
  1.读入一棵树，由于分叉树未知，不妨用vector<int>的数组来存放每个结点的孩子。
  2.利用深搜递归地进行编号，用一个二元组pair<int,int>来存放左值和右值
  3.构造线段树，其根节点的范围为[1,n]，对应刚才的编号。要注意，起始阶段每个结点的值均为1，故每个区间的值都为其长度。
  4.编写修改函数。修改是递归的，每一次修改需要下一层修改的返回值来确定是增加了苹果还是减少了苹果。
  5.编写查询函数。标准的线段树操作即可。

代码如下

#include<iostream>
#include<vector>
#include<algorithm>
#include<cstring>

#define MAXN 100005
#define MAXM 100005

using namespace std;

int N, M;
vector<int> child[MAXN];    //存储每个节点的孩子，用于dfs
int id[MAXN];               //存储每个节点在dfs过程中的编号id，从1开始，以下的编号全部是指dfs的编号
int range[MAXN];            //存储每个节点所有的孩子中，最大的编号,这样[i,range[i]]就是要求和的区间范围
int x[MAXN];                //树状数组，可以用来求第1个至第i个节点的苹果数之和
int apples[MAXN];           //用来存储第1至第i个节点是否有苹果

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

int sum(int k) {
    int ans = 0;
    while (k) {
        ans += x[k];
        k -= lowbit(k);
    }
    return ans;
}

void change(int k) {        //根据apples[k]中的值调整树状数组
    int bias = -1;
    if (apples[k] == 1) bias = 1;
    while (k<=N) {          //向后调整树状数组
        x[k] += bias;
        k += lowbit(k);
    }
}

void init() {
    memset(apples, 0, sizeof(apples));
    memset(x, 0, sizeof(x));
    for (int i = 1; i <= N; i++) {
        apples[i] = 1;
        child[i].clear();
    }
    for (int i = 1; i <= N; i++) {
        change(i);
    }
}

int dfsCnt = 0;
void dfs(int k) {
    dfsCnt++;
    id[k] = dfsCnt;
    typedef vector<int>::iterator it;
    for (it i = child[k].begin(); i!= child[k].end(); ++i) {
        dfs(*i);
    }
    range[id[k]] = dfsCnt;
}

int main() {
    cin >> N;
    init();
    for (int i = 1; i <= N - 1; i++) {
        int father, son;
        cin >> father >> son;
        child[father].push_back(son);
    }
    dfs(1);
    cin >> M;
    char op;
    int k;
    for (int i = 1; i <= M; i++) {
        cin >> op >> k;
        if (op == 'Q') {
            int l = id[k];
            int r = range[l];
            cout << (sum(r) - sum(l - 1)) << endl;  
        }
        else if (op == 'C') {
            if (k == 2) {
                int stophere = 2;
            }
            apples[id[k]] = !apples[id[k]];
            change(id[k]);  
        }
    }
}