(Sky OJ) 146. Day 3 PH. 樹樹鏈鏈剖剖分分 [樹鏈剖分]

https://pc2.tfcis.org/sky/index.php/problem/view/146/

#include <iostream>
#include <stdio.h>
#include <vector>
#include <algorithm>
#include <cstring>
#include <queue>
using namespace std;

typedef long long LL;
const int MAX_N = 2e5 +6;
const int MAX_P = 19;

vector<int> edg[MAX_N];
int depth[MAX_N];
int p[MAX_N];
int lca[MAX_P][MAX_N];
int pin[MAX_N],pout[MAX_N];
int sz[MAX_N];
int stamp;

void dfs(int id,int p,int cur_depth) {
    lca[0][id] = p;
    depth[id] = cur_depth;
    pin[id] = stamp++;
    sz[id] = 1;
    for (int i:edg[id])  {
        if (i!=p) {
            dfs(i,id,cur_depth+1);
            sz[id] += sz[i];
        }
    }
    pout[id] = stamp++;
}

bool is_anc(int parent,int son ){
    return pin[parent] <= pin[son] && pout[son] <= pout[parent];
}

int get_lca(int a,int b) {
    if (depth[a] > depth[b]) swap(a,b);
    if (is_anc(a,b)) return a;
    int bb=b;
    for (int i=MAX_P-1;i>=0;i--) {
        if (!is_anc(lca[i][b],a)) b=lca[i][b];
    }
    b=lca[0][b];
    return b;
}

LL seg[4*MAX_N];
int tag[4*MAX_N];

void push(int id,int L,int R) {
    if (L==R) tag[id]=0;
    if (tag[id]==0) return;
    tag[id*2] += tag[id];
    tag[id*2+1] += tag[id];
    int mid=(L+R)>>1;
    seg[id*2] += tag[id]*(mid-L+1);
    seg[id*2+1] += tag[id]*(R-mid);
    tag[id]=0;
    return;
}

void modify(int id,int L,int R,int l,int r,int val) {
    if (l>R || L>r) return;
    else if (l<=L && R<=r) {
        tag[id] += val;
        seg[id] += (R-L+1)*val;
        return;
    }
    push(id,L,R);
    int mid=(L+R)>>1;
    if (mid + 1 > r) modify(id*2,L,mid,l,r,val);
    else if (l > mid) modify(id*2+1,mid+1,R,l,r,val);
    else {
        modify(id*2,L,mid,l,r,val);
        modify(id*2+1,mid+1,R,l,r,val);
    }
    seg[id] = seg[id*2] + seg[id*2+1];
}

LL query(int id,int L,int R,int l,int r){
    if (l>R || L>r) return 0;
    else if (l<=L && R<=r) return seg[id];
    push(id,L,R);
    int mid=(L+R)>>1;
    if (mid + 1 > r) return query(id*2,L,mid,l,r);
    else if (l > mid) return query(id*2+1,mid+1,R,l,r);
    else {
        return query(id*2,L,mid,l,r) + query(id*2+1,mid+1,R,l,r);
    }
}

#define SZ(x) ((int)(x).size())

int head[MAX_N];
int tin[MAX_N];

int n,m;

void dfs2(int id,int p,bool nnew) {
    sort(edg[id].begin(),edg[id].end(),[](const int &a,const int &b) {
        return sz[a] >sz[b];
    });
    if (nnew) {
        head[id] = id;
    }
    else {
        head[id] = head[p];
    }
    tin[id] = ++stamp;
    bool gett=false;
    for (int ii=0;SZ(edg[id])>ii;ii++) {
        int i=edg[id][ii];
        if (i==p) continue;
        dfs2(i,id,gett);
        gett=true;
    }
}

LL Q(int son,int parent) {
    if (son == parent) return 0;
    LL ret=0;
    while (head[son] != head[parent]) {
        //son --> head[son]
        ret += query(1,1,n,tin[head[son]],tin[son]);
        son = head[son];
        son = lca[0][son];
    }
    if (son == parent) return ret+0;
    else return ret+query(1,1,n,tin[parent]+1,tin[son]);
}

void change(int son,int parent) {
    if (son == parent) return;
    while (head[son] != head[parent]) {
        //son --> head[son]
        modify(1,1,n,tin[head[son]],tin[son],1);
        son = head[son];
        son = lca[0][son];
    }
    if (son == parent) return;
    else modify(1,1,n,tin[parent]+1,tin[son],1);
}

int main () {
    scanf("%d %d",&n,&m);
    for (int i=1;n>i;i++) {
        int a,b;
        scanf("%d %d",&a,&b);
        edg[a].push_back(b);
        edg[b].push_back(a);
    }
    stamp=0;
    dfs(1,1,1);
    for (int i=1;MAX_P>i;i++) {
        for (int j=1;n>=j;j++) {
            lca[i][j] = lca[i-1][lca[i-1][j]];
        }
    }
    stamp=0;
    dfs2(1,1,1);
    for (int i=1;m>=i;i++) {
        getchar();
        char c;
        int a,b;
        scanf("%c %d %d",&c,&a,&b);
        int lca = get_lca(a,b);
        if (c=='P') {
            change(a,lca);
            change(b,lca);
        }
        else if (c=='Q') {
            printf("%lld\n",Q(a,lca) + Q(b,lca));
        }
    }
}