HDU 5893 List wants to travel（树链剖分+线段树）

Description

给出一棵$n$个点的树，每条边有颜色，$m$次操作，操作分两种

$Change\ a\ b\ c:$把从$a$点到$b$点路径上所有边的颜色变成$c$

$Query\ a\ b:$查询$a$点到$b$点路径上的边有几段连续颜色相同

Input

第一行两个整数$n,m$表示点数和操作数，之后$n-1$行每行三个整数$u,v,c$表示$u,v$之间有一条树边颜色为$c$，之后$m$行每行一个查询$(1\le n\le 40000,1\le m\le 50000,1\le c\le 10^5)$

Output

对于每组查询操作，输出查询结果

Sample Input

9 3
1 2 2
2 3 1
1 7 2
1 4 2
3 5 2
3 6 1
5 8 2
5 9 3
Query 1 8
Change 2 6 3
Query 1 6

Sample Output

3
2

Solution

以$1$为根节点对整棵树树链剖分，把边的颜色放到深度较深的点上变为点的颜色，由$dfs$序把$a$到$b$的路径修改查询变成对若干线段的查询修改，进而问题转化为区间更新和区间段数查询，线段树记录区间段数，区间左右端点颜色以及区间是否全部为同一颜色即可

Code

#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
#define INF -0x3f3f3f3f
#define maxn 111111 
struct Edge
{
    int to,next;
}edge[2*maxn];
int head[maxn],tot;
int idx;//时间戳,用于得到dfs序 
int size[maxn];//size[i]表示以i节点为根节点的子树中节点个数 
int fa[maxn];// fa[i]表示i节点的父亲节点 
int son[maxn];//son[i]表示i节点的重儿子节点(如果没有重儿子则son[i]=i) 
int dep[maxn];//dep[i]表示i节点在树中的深度 
int top[maxn];//top[i]表示i节点所处重链中深度最低的节点 
int id[maxn];//id[i]表示i节点的dfs序 
int pos[maxn];//pos[i]表示dfs序为i的节点 
void init()
{
    idx=tot=0;
    dep[1]=fa[1]=size[0]=0;//默认1节点为根节点,初始化其父亲节点为0,深度为0 
    memset(son,0,sizeof(son));
    memset(head,-1,sizeof(head));
}
void add(int u,int v)
{
    edge[tot].to=v;
    edge[tot].next=head[u];
    head[u]=tot++;
}
void dfs1(int u)//得到size,fa,dep,son
{
    size[u]=1;
    for(int i=head[u];~i;i=edge[i].next)
    {
        int v=edge[i].to;
        if(v!=fa[u]) 
        {
            fa[v]=u;//v的父亲节点是u 
            dep[v]=dep[u]+1;//儿子的深度等于父亲的深度加一 
            dfs1(v);//深搜 
            size[u]+=size[v];//父亲节点为根节点的子树节点个数等于各儿子节点为根节点的子树节点个数之和加一(父亲节点本身) 
            if(size[son[u]]<size[v]) son[u]=v;//更新重儿子 
        }
    }
}
void dfs2(int u,int topu)//得到top,id,pos,l,r 
{
    top[u]=topu;
    id[u]=++idx;//得到u节点的dfs序 
    pos[idx]=u;//记录这个dfs序对应的节点 
    if(son[u]) dfs2(son[u],top[u]);//有重儿子首先深搜重儿子 
    for(int i=head[u];~i;i=edge[i].next)
    {
        int v=edge[i].to;
        if(v!=fa[u]&&v!=son[u]) dfs2(v,v);//深搜所有儿子节点 
    }
}
int n,m,e[maxn][3],a[maxn];
#define ls (t<<1)
#define rs ((t<<1)|1)
int Same[maxn<<2],Num[maxn<<2],LL[maxn<<2],RR[maxn<<2];
void push_up(int t)
{
    LL[t]=LL[ls],RR[t]=RR[rs];
    Num[t]=Num[ls]+Num[rs];
    if(RR[ls]==LL[rs])Num[t]--;
    if(Same[ls]==Same[rs])Same[t]=Same[ls];
}
void push_down(int l,int r,int t)
{
    if(Same[t]!=INF)
    {
        Same[ls]=LL[ls]=RR[ls]=Same[rs]=LL[rs]=RR[rs]=Same[t];
        Num[ls]=Num[rs]=1;
        Same[t]=INF;
    }
}
void build(int l,int r,int t)
{
    Same[t]=LL[t]=RR[t]=INF,Num[t]=0;
    if(l==r)
    {
        Same[t]=LL[t]=RR[t]=a[l],Num[t]=1;
        return ;
    }
    int mid=(l+r)>>1;
    build(l,mid,ls),build(mid+1,r,rs);
    push_up(t);
}
void update(int L,int R,int l,int r,int t,int v)
{
    if(L==l&&r==R)
    {
        Same[t]=LL[t]=RR[t]=v,Num[t]=1;
        return ;
    }
    push_down(l,r,t);
    int mid=(l+r)>>1;
    if(R<=mid)update(L,R,l,mid,ls,v);
    else if(L>mid)update(L,R,mid+1,r,rs,v);
    else update(L,mid,l,mid,ls,v),update(mid+1,R,mid+1,r,rs,v);
    push_up(t);
}
int query(int L,int R,int l,int r,int t)
{
    if(L==l&&r==R)return Num[t];
    push_down(l,r,t);
    int mid=(l+r)>>1;
    if(R<=mid)return query(L,R,l,mid,ls);
    else if(L>mid)return query(L,R,mid+1,r,rs);
    else
    {
        int ans=query(L,mid,l,mid,ls)+query(mid+1,R,mid+1,r,rs);
        if(RR[ls]==LL[rs])ans--;
        return ans;
    }
    push_up(t);
}
int ask(int x,int l,int r,int t)
{
    if(l==r)return LL[t];
    push_down(l,r,t);
    int mid=(l+r)>>1;
    if(x<=mid)return ask(x,l,mid,ls);
    else return ask(x,mid+1,r,rs);
}
void Update(int u,int v,int vv)
{
    int top1=top[u],top2=top[v];
    while(top1!=top2)
    {
        if(dep[top1]<dep[top2])
        {
            swap(top1,top2);
            swap(u,v);
        }
        update(id[top1],id[u],1,n,1,vv);
        u=fa[top1];
        top1=top[u];
    }
    if(dep[u]>dep[v]) swap(u,v);
    if(u!=v)update(id[son[u]],id[v],1,n,1,vv);
}
int Query(int u,int v)
{
    int top1=top[u],top2=top[v],ans=0;
    while(top1!=top2)
    {
        if(fa[u]==fa[v])
        {
            ans+=query(id[u],id[u],1,n,1)+query(id[v],id[v],1,n,1);
            if(ask(id[u],1,n,1)==ask(id[v],1,n,1))ans--;
            return ans;
        }
        if(dep[top1]<dep[top2])
        {
            swap(top1,top2);
            swap(u,v);
        }
        ans+=query(id[top1],id[u],1,n,1);
        if(top[fa[top1]]==top2)
        {
            if(dep[fa[top1]]<=dep[v])
            {
                if(fa[top1]==v)return ans;
                if(ask(id[top1],1,n,1)==ask(id[son[fa[top1]]],1,n,1))ans--;
                u=fa[top1];
                break;
            }
        }
        if(ask(id[top1],1,n,1)==ask(id[fa[top1]],1,n,1))ans--;
        u=fa[top1];
        top1=top[u];
    }
    if(dep[u]>dep[v])swap(u,v);
    if(u!=v)ans+=query(id[son[u]],id[v],1,n,1);
    return ans;
}
int main()
{
    while(~scanf("%d%d",&n,&m))
    {
        init();
        for(int i=1;i<n;i++)
        {
            int u,v,w;
            scanf("%d%d%d",&u,&v,&w);
            e[i][0]=u,e[i][1]=v,e[i][2]=w;
            add(u,v),add(v,u);
        }
        dfs1(1);
        dfs2(1,1); 
        for(int i=1;i<n;i++)
        {
            int u=e[i][0],v=e[i][1],w=e[i][2];
            if(dep[u]<dep[v])a[id[v]]=w;
            else a[id[u]]=w;
        }
        build(1,n,1);
        while(m--)
        {
            char op[11];
            int u,v,vv;
            scanf("%s%d%d",op,&u,&v);
            if(op[0]=='Q')
            {
                if(u==v)printf("0\n");
                else printf("%d\n",Query(u,v));
            }
            else 
            {
                scanf("%d",&vv);
                Update(u,v,vv);
            }
        }
    }
    return 0;
}