hdu 5893 List wants to travel 树链剖分求区间段数

用到树链剖分和线段树合并，线段数存储段数即需要包含左端点颜色右端点颜色和段数和，但是难点在树上的查询的时候会牵扯到 不同链之间的交接处是否重色，参考了别人的写法但是代码写的还是冗余度很高，好累，先放着吧，看以后有没有时间来精简下。

#pragma comment(linker, "/STACK:1024000000,1024000000")
#include <stdio.h>
#include <string.h>
#include <iostream>
#include <algorithm>
#include <vector>
#include <queue>
#include <set>
#include <map>
#include <string>
#include <math.h>
#include <stdlib.h>
using namespace std;

const int MAXN = 100010;
int add[MAXN*3];
struct Edge
{
    int to,next;
}edge[MAXN*2];
int head[MAXN],tot;
int top[MAXN]; //top[v] 表示v所在的重链的顶端节点
int fa[MAXN]; //父亲节点
int deep[MAXN];//深度
int num[MAXN]; //num[v]表示以v为根的子树的节点数
int p[MAXN]; //p[v]表示v在线段树中的位置
int fp[MAXN];//和p数组相反
int son[MAXN];//重儿子
int pos;
void init()
{
    tot = 0;
    memset(head,-1,sizeof(head));
    pos = 0;
    memset(son,-1,sizeof(son));
    memset(add,0,sizeof(add));
}
void addedge(int u,int v)
{
    edge[tot].to = v;edge[tot].next = head[u]; head[u] = tot++;
}
void dfs1(int u,int pre,int d) //第一遍dfs求出fa,deep,num,son
{
    deep[u] = d;
    fa[u] = pre;
    num[u] = 1;
    for(int i = head[u];i != -1;i = edge[i].next)
    {
        int v = edge[i].to;
        if(v != pre)
        {
            dfs1(v,u,d+1);
            num[u] += num[v];
            if(son[u] == -1 || num[v] > num[son[u]])
                son[u] = v;
        }
    }
}
void getpos(int u,int sp)
{
    top[u] = sp;
    p[u] = pos++;
    fp[p[u]] = u;
    if(son[u] == -1) return;
    getpos(son[u],sp);
    for(int i = head[u]; i != -1 ; i = edge[i].next)
    {
        int v = edge[i].to;
        if(v != son[u] && v != fa[u]) getpos(v,v);
    }
}

struct Node
{
    int lc,rc,c;
    int l,r;
    int sum;
}segTree[MAXN*3];

void push_up(int i)
{
    segTree[i].sum = segTree[i<<1].sum + segTree[(i<<1)|1].sum-(segTree[i<<1].rc==segTree[(i<<1)|1].lc);
    segTree[i].lc=segTree[i<<1].lc;
    segTree[i].rc=segTree[(i<<1)|1].rc;
    //segTree[i].Max = max(segTree[i<<1].Max,segTree[(i<<1)|1].Max);
}
void push_down(int i)
{
    if(add[i]!=0)
    {
        segTree[i*2].c=segTree[i*2].lc=segTree[i*2].rc=add[i];
        segTree[i*2+1].c=segTree[i*2+1].lc=segTree[i*2+1].rc=add[i];
        add[i*2]=add[i];
        add[i*2+1]=add[i];
        segTree[i<<1].sum=1;
        segTree[(i<<1)|1].sum=1;
        add[i]=0;
    }
}
void build(int i,int l,int r)
{
    segTree[i].l = l;
    segTree[i].r = r;
    segTree[i].lc=segTree[i].rc=segTree[i].c=0;
    if(l == r)
    {
        segTree[i].sum = 0;
        return ;
    }
    int mid = (l + r)/2;
    build(i<<1,l,mid);
    build((i<<1)|1,mid+1,r);
    segTree[i].sum = segTree[i<<1].sum + segTree[(i<<1)|1].sum;
}

void update(int u,int l,int r,int v)
{
    if(segTree[u].l == l && segTree[u].r == r)
    {
        segTree[u].c = segTree[u].lc=segTree[u].rc=v;
        segTree[u].sum=1;
        add[u]=v;
        return;
    }
    push_down(u);
    int mid = (segTree[u].l + segTree[u].r) >> 1;
    if(r <= mid)
        update(u*2, l, r, v);
    else if(l > mid)
        update(u*2+1, l, r, v);
    else
    {
        update(u*2, l, mid, v);
        update(u*2+1, mid + 1, r, v);
    }
    push_up(u);
}
int querySum(int i,int l,int r) //查询线段树[l,r]区间的和
{
    if(segTree[i].l == l && segTree[i].r == r)
        return segTree[i].sum;
        push_down(i);
    int mid = (segTree[i].l + segTree[i].r)/2;
    if(r <= mid)return querySum(i<<1,l,r);
    else if(l > mid)return querySum((i<<1)|1,l,r);
    else return querySum(i<<1,l,mid) + querySum((i<<1)|1,mid+1,r) -(segTree[i<<1].rc==segTree[i<<1|1].lc);
}
int query_single(int i, int l, int r, int x)    //查询单点颜色
    {
        if(l ==r )  return add[i];
        push_down(i);              //同update
        int m=(l+r)/2;
        if(x<=m) return query_single(i<<1,l,m,x);
        return query_single(i<<1|1,m+1,r,x);
    }
int findsum(int u,int v)//查询u->v链的sum
    {
        int f1=top[u],f2=top[v];
        int tmp=0;
        int lastu=-1,col,lastv=-1;
        while(f1!=f2)
        {
            if (deep[f1]<deep[f2])
            {
                swap(f1,f2);
                swap(u,v);
                swap(lastu,lastv);
            }
            tmp+=querySum(1,p[f1],p[u]);

            col=query_single(1,0,pos-1,p[u]);
            if (col==lastu) tmp--;
            lastu=query_single(1,0,pos-1,p[f1]);

            u=fa[f1],f1=top[u];
        }
        if (u!=v)
        {
            if (deep[u]>deep[v] ) swap(u,v),swap(lastu,lastv);

            tmp+=querySum(1,p[son[u]],p[v]);

            col=query_single(1,0,pos-1,p[v]);
            if (col==lastv) tmp--;
            col=query_single(1,0,pos-1,p[son[u]]);
            if (col==lastu) tmp--;
        }
        else if (lastu==lastv) tmp--;
        return  tmp; 
    }
void change(int u,int v,int val)
{
    int f1 = top[u], f2 = top[v];
    int tmp = 0;
    while(f1 != f2)
    {
        if(deep[f1] < deep[f2])
        {
            swap(f1,f2);
            swap(u,v);
        }
        update(1,p[f1],p[u],val);
        u = fa[f1];
        f1 = top[u];
    }
    if(deep[u] > deep[v]) swap(u,v);
     if (u!=v)
        update(1,p[son[u]],p[v],val);
}
int e[MAXN][3];
int main()
{
    int M,P,n;
    while(scanf("%d%d",&n,&M) !=EOF)
    {
        int u,v;
        int C1,C2,K;
        char op[10];
        init();
        for(int i = 0;i <n-1;i++)
        {
            scanf("%d%d%d",&e[i][0],&e[i][1],&e[i][2]);
            addedge(e[i][0],e[i][1]);
            addedge(e[i][1],e[i][0]);
        }
        dfs1(1,0,0);
        getpos(1,1);
        build(1,0,pos-1);
        for(int i=0;i<n-1;i++)
        {
            if(deep[e[i][0]]>deep[e[i][1]])
                swap(e[i][0],e[i][1]);
            update(1,p[e[i][1]],p[e[i][1]],e[i][2]);
        }
        while(M--)
        {
            scanf("%s",op);
            if(op[0] == 'Q')
            {
                scanf("%d%d",&u,&v);
                //printf("%d\n",sum(p[u]));
                cout<<findsum(u,v)<<endl;
            }
            else
            {
                scanf("%d%d%d",&C1,&C2,&K);
                change(C1,C2,K);
            }
        }
    }
    return 0;
}