树的统计Count HYSBZ

树的统计Count

HYSBZ - 1036

高级数据结构p329

感叹一句分块大法吼啊...

#include <cstdio>
#include <cstring>
#include <iostream>
#include <cmath>
#include <algorithm>
using namespace std;
const int MAXN = 310000;
const int inf = 0x3f3f3f3f;

//struct Edge{
//    int v, nxt;
//}e1[MAXN<<1], e2[MAXN<<1];
//int head1[MAXN], head2[MAXN];
//int cnt1, cnt2;
//void init(int &cnt, int *&head){
//    memset(head, -1, sizeof(head));
//    cnt = 0;    
//}
//void add(int u, int v, int *&head, Edge *&e, int &cnt){
//    e[cnt++] = Edge{v, head[u]};
//    head[u] = cnt++;
//}

struct Edge{
    int cnt, head[MAXN<<1], nxt[MAXN<<1], to[MAXN<<1];
    void init(){
        memset(head, -1, sizeof(head));
        cnt = 0;
    }
    void add(int u, int v){
        to[cnt] = v;
        nxt[cnt] = head[u];
        head[u] = cnt++;
    }
}e1, e2;
int val[MAXN], dep[MAXN], f[MAXN], bkrt[MAXN];
int sz[MAXN], limit;
int VSum[MAXN], VMax[MAXN];

void buildBlock(int u, int fa, int d){
    dep[u] = d;
    f[u] = fa;
    int curB = bkrt[u];
    for(int i = e1.head[u]; ~i; i = e1.nxt[i]){
        int v = e1.to[i];
        if(v != fa){
            if(sz[curB] + 1 < limit){
                e2.add(u, v);
                bkrt[v] = curB;
                sz[curB]++;
            }
            buildBlock(v, u, d + 1);
        }
    }
}

void InitData(int u, int sumval, int maxval){
    sumval += val[u]; maxval = max(maxval, val[u]);
    VSum[u] = sumval; VMax[u] = maxval;
    for(int i = e2.head[u]; ~i; i = e2.nxt[i]){
        InitData(e2.to[i], sumval, maxval);
    }
}
void update(int u, int data){
    val[u] = data;
    if(u == bkrt[u]) InitData(u, 0, -inf);
    else InitData(u, VSum[f[u]], VMax[f[u]]);
}
pair<int, int> query(int a, int b){
    pair<int, int> ans(0, -inf);
    while(a != b){
        if(dep[a] < dep[b]) swap(a, b);
        if(bkrt[a] == bkrt[b]){
            ans.first += val[a];
            ans.second = max(ans.second, val[a]);
            a = f[a];
        }else{
            if(dep[bkrt[a]] < dep[bkrt[b]]) swap(a, b);
            ans.first += VSum[a];
            ans.second = max(ans.second, VMax[a]);
            a = f[bkrt[a]];
        }
    }
    ans.first += val[a];
    ans.second = max(ans.second, val[a]);
    return ans;
}

int main(){
    int n;
    int u, v;
    scanf("%d", &n);
    limit = sqrt(n) + 1;
    e1.init(); e2.init();
    for(int i = 1; i < n; i++){
        scanf("%d %d", &u, &v);
        e1.add(u, v);
        e1.add(v, u);
    }
    for(int i = 1; i <= n; i++){
        scanf("%d", &val[i]);
        bkrt[i] = i;
    }
    memset(sz, 0, sizeof(sz));
    buildBlock(1, 0, 0);
    for(int i = 1; i <= n; i++){
        if(bkrt[i] == i){
            InitData(i, 0, -inf);
        }
    }
    int m;
    char op[20];
    scanf("%d", &m);
    for(int i = 0; i < m; i++){
        scanf("%s %d %d", op, &u, &v);
        pair<int, int> ans;
        if(op[1] == 'M'){
            ans = query(u, v);
            printf("%d
", ans.second);
        }else if(op[1] == 'S'){
            ans = query(u, v);
            printf("%d
", ans.first);
        }else{
            update(u, v);
        }
    }
    return 0;
}

虽然还不是很理解，，先码下吧