HDU 5893 List wants to travel 【线段树+树链剖分】

#include<bits/stdc++.h>
using namespace std;
typedef vector<int> VI;
typedef unsigned long long ull;
typedef long long ll;
typedef double db;
const ll mod=1000000007;
const int N=2e5+10;
const int inf=0x7fffffff;
struct edge
{
    int a,b,w;
} g[2*N];
vector<int>v[N];
int fa[N],dep[N],sz[N],son[N],w[N],vis[N],top[N],tot;

void dfs1(int x,int f)
{
    int y,mx=-1;
    vis[x]=1;
    fa[x]=f;
    sz[x]=1;
    son[x]=0;
    dep[x]=dep[f]+1;
    for(int i=0; i<v[x].size(); i++)
    {
        y=v[x][i];
        if(vis[y]==1) continue;
        dfs1(y,x);
        sz[x]+=sz[y];
        if(sz[y]>mx)
        {
            mx=sz[y];
            son[x]=y;
        }
    }
}

void dfs2(int x,int tp)
{
    top[x]=tp;
    if(son[x]!=0)
    {
        w[x]=tot++;
        dfs2(son[x],tp);
    }
    else
    {
        w[x]=tot++;
        return ;
    }
    int y;
    for(int i=0; i<v[x].size(); i++)
    {
        y=v[x][i];
        if(y!=son[x]&&y!=fa[x])
            dfs2(y,y);
    }
}

int mx[4*N],L[4*N],R[4*N],f[4*N];

void pushup(int o)
{
    L[o]=L[o*2];
    R[o]=R[o*2+1];
    if(R[o*2]==L[o*2+1])
    {
        mx[o]=mx[o*2]+mx[o*2+1]-1;
    }
    else
    {
        mx[o]=mx[o*2]+mx[o*2+1];
    }
}

void pushdw(int o)
{
    if(f[o]!=-1)
    {
        f[o*2]=f[o*2+1]=f[o];
        L[o*2]=R[o*2]=L[o*2+1]=R[o*2+1]=f[o];
        mx[o*2]=mx[o*2+1]=1;
        f[o]=-1;
    }
}

void build(int o,int l,int r)
{
    L[o]=R[o]=mx[o]=0;
    f[o]=-1;
    if(l==r) return ;
    int m=(l+r)/2;
    build(o*2,l,m);
    build(o*2+1,m+1,r);
}

int query(int o,int l,int r,int x,int y)
{
    if(x<=l&&r<=y)
    {
        return mx[o];
    }
    else
    {
        pushdw(o);
        int m=(l+r)/2;
        if(y<=m)
        {
            return query(o*2,l,m,x,y);
        }
        else if(x>m)
        {
            return query(o*2+1,m+1,r,x,y);
        }
        else
        {
            int a1=query(o*2,l,m,x,y);
            int a2=query(o*2+1,m+1,r,x,y);
            int ret=a1+a2;
            if(R[o*2]==L[o*2+1]) ret--;
            return ret;
        }
        pushup(o);
    }
}

int query2(int o,int l,int r,int x)
{
    if(f[o]!=-1)
    {
        return f[o];
    }
    else if(l==r)
    {
        return L[o];
    }
    else
    {
        pushdw(o);
        int m=(l+r)/2;
        if(x<=m)
            return query2(o*2,l,m,x);
        else if(x>m)
            return query2(o*2+1,m+1,r,x);
        pushup(o);
    }
}
void update(int o,int l,int r,int x,int y,int v)
{
    pushdw(o);
    if(x<=l&&r<=y)
    {
        f[o]=v;
        mx[o]=1;
        L[o]=v;
        R[o]=v;
    }
    else
    {
        int m=(l+r)/2;
        if(x<=m) update(o*2,l,m,x,y,v);
        if(y>m) update(o*2+1,m+1,r,x,y,v);
        pushup(o);
    }
}

int Find(int x,int y,int n)
{
    int f1=top[x],f2=top[y],p1=-1,p2=-1,c1,c2;
    int ans=0;
    while(f1!=f2)
    {
        if(dep[f1]<dep[f2])
        {
            swap(f1,f2);
            swap(x,y);
            swap(p1,p2);
            swap(c1,c2);
        }
        ans+=query(1,1,n,w[f1],w[x]);
        int now=query2(1,1,n,w[x]);
        if(p1!=-1&&fa[p1]==x)
        {
            if(c1==now) ans--;
        }
        p1=f1;
        c1=query2(1,1,n,w[f1]);
        x=fa[f1];
        f1=top[x];
    }
    if(x==y)
    {
        if(p1!=-1&&p2!=-1&&c1==c2) ans--;
        return ans;
    }
    if(dep[x]>dep[y])
    {
        swap(x,y);
        swap(p1,p2);
        swap(c1,c2);
    }
    ans+=query(1,1,n,w[son[x]],w[y]);
    //printf("%d %d %d %d\n",p1,p2,son[x],y);
    if(p2!=-1&&fa[p2]==y)
    {
        if(c2==query2(1,1,n,w[y])) ans--;
    }
    if(p1!=-1&&fa[p1]==x)
    {
        if(c1==query2(1,1,n,w[son[x]])) ans--;
    }
    return ans;
}

void Change(int x,int y,int c,int n)
{
    int f1=top[x],f2=top[y];
    while(f1!=f2)
    {
        if(dep[f1]<dep[f2])
        {
            swap(f1,f2);
            swap(x,y);
        }
        update(1,1,n,w[f1],w[x],c);
        x=fa[f1];
        f1=top[x];
    }
    if(x==y) return ;
    if(dep[x]>dep[y])
        swap(x,y);
    update(1,1,n,w[son[x]],w[y],c);
}

int main()
{
    //freopen("in.txt","r",stdin);
    //freopen("out.txt","w",stdout);
    int _,n,m,a,b,c;
    char s[100];
    while(~scanf("%d%d",&n,&m))
    {
        for(int i=1; i<=n; i++)
            v[i].clear();
        for(int i=0; i<n-1; i++)
        {
            scanf("%d%d%d",&a,&b,&c);
            g[i].a=a;
            g[i].b=b;
            g[i].w=c;
            v[a].push_back(b);
            v[b].push_back(a);
        }
        memset(vis,0,sizeof(vis));
        memset(son,0,sizeof(son));
        tot=1;
        dep[0]=0;
        dfs1(1,0);
        dfs2(1,1);
        build(1,1,tot-1);
        for(int i=0; i<n-1; i++)
        {
            a=g[i].a;
            b=g[i].b;
            if(dep[a]>dep[b])
                swap(g[i].a,g[i].b);
            update(1,1,n,w[g[i].b],w[g[i].b],g[i].w);
        }
        for(int i=0; i<m; i++)
        {
            scanf("%s",s);
            scanf("%d%d",&a,&b);
            if(s[0]=='C')
            {
                scanf("%d",&c);
                Change(a,b,c,n);
            }
            else
            {
                printf("%d\n",Find(a,b,n));
            }
        }
    }
    return 0;
}