HDU 5893 List wants to travel 树链剖分求区间不同段个数

传送门：HDU5893

题意：给出一棵树和2种操作 
1.修改u到v的路径上的边的颜色为c 
2.查询u到v的路径上有多少段颜色（连续相同的颜色算一段）

思路：明显是树链剖分的题，难点在于维护颜色的连续性，在线段树中还好说，基本的线段树套路，左右区间邻接处比较一下就好了，关键是在询问过程中两边不断上移的过程，需要记录上一次询问的链的顶端颜色，和这一次询问链的底端颜色比较，看看是否能连成一段，由于是两边交替上移，因此两边的颜色都要记录，最后（u，v）合到一条链上去以后还要注意判断一下和原来的两条链能否颜色相接。

思路很容易建立，代码蒟蒻debug了4个小时（还是在和题解对拍的情况下）。。

代码：

#include<bits/stdc++.h>
#define ll long long
#define inf 0x3f3f3f3f
#define lson l, mid, rt << 1
#define rson mid + 1, r, rt << 1 | 1
using namespace std;
typedef pair<int,int> P;
const int MAXN = 50010;
int fa[MAXN], top[MAXN], pos[MAXN], son[MAXN], sz[MAXN], dep[MAXN];
int tid;
vector<int> mp[MAXN];
int color[MAXN];
struct node{
    int l, r;
    int lc, rc, lazy;
    int cnt;
}tree[MAXN << 2];
struct edge{
    int u, v, w;
}g[MAXN];
void init(int n)
{
    tid = 0;
    for(int i = 0; i <= n; i++) mp[i].clear();
    memset(son, -1, sizeof(son));
}
void dfs1(int u)
{
    sz[u] = 1;
    int v;
    for(int i = 0; i < mp[u].size(); i++)
    {
        v = mp[u][i];
        if(v == fa[u]) continue;
        fa[v] = u; dep[v] = dep[u] + 1;
        dfs1(v);
        sz[u] += sz[v];
        if(son[u] == -1 || sz[son[u]] < sz[v])
        son[u] = v;
    }
}
void dfs2(int u, int head)
{
    top[u] = head;
    pos[u] = ++tid;
    if(son[u] == -1) return ;
    dfs2(son[u], head);
    int v;
    for(int i = 0; i < mp[u].size(); i++)
    {
        v = mp[u][i];
        if(v == fa[u] || v == son[u]) continue;
        dfs2(v, v);
    }
}
void push_up(int rt)
{
    tree[rt].cnt = tree[rt << 1].cnt + tree[rt << 1 | 1].cnt;
    if(tree[rt << 1].rc == tree[rt << 1 | 1].lc)
    tree[rt].cnt--;
    tree[rt].lc = tree[rt << 1].lc;
    tree[rt].rc = tree[rt << 1 | 1].rc;
}
void push_down(int rt)
{
    if(tree[rt].lazy == -1) return ;
    tree[rt << 1].lc = tree[rt << 1].rc = tree[rt].lazy;
    tree[rt << 1 | 1].lc = tree[rt << 1 | 1].rc = tree[rt].lazy;
    tree[rt << 1].cnt = tree[rt << 1 | 1].cnt = 1;
    tree[rt << 1].lazy = tree[rt << 1 | 1].lazy = tree[rt].lazy;
    tree[rt].lazy = -1;
}
void build(int l, int r, int rt)
{
    tree[rt].l = l;
    tree[rt].r = r;
    tree[rt].lazy = -1;
    if(l == r)
    {
        tree[rt].cnt = 1;
        tree[rt].lc = tree[rt].rc = color[l];
        return ;
    }
    int mid = (l + r) >> 1;
    build(lson); build(rson);
    push_up(rt);
}
void update(int rt, int L, int R, int x)
{
    if(L <= tree[rt].l && tree[rt].r <= R)
    {
        tree[rt].lazy = tree[rt].lc = tree[rt].rc = x;
        tree[rt].cnt = 1;
        return ;
    }
    push_down(rt);
    int mid = (tree[rt].l + tree[rt].r) >> 1;
    if(L <= mid)
    update(rt << 1, L, R, x);
    if(R > mid)
    update(rt << 1 | 1, L, R, x);
    push_up(rt);
}
int query(int rt, int L, int R, int &lc, int &rc)
{
    if(L <= tree[rt].l && tree[rt].r <= R)
    {
        if(L == tree[rt].l)
        lc = tree[rt].lc; 
        if(R == tree[rt].r)
        rc = tree[rt].rc;
        return tree[rt].cnt;
    }
    push_down(rt);
    int mid = (tree[rt].l + tree[rt].r) >> 1, ans = 0;
    if(R <= mid)
	ans = query(rt << 1, L, R, lc, rc);
    else if(L > mid)
	ans = query(rt << 1 | 1, L, R, lc, rc);
    else
    {
        int tl, tr;
        ans = query(rt << 1, L, R, lc, tr) + query(rt << 1 | 1, L, R, tl, rc);
        if(tree[rt << 1].rc == tree[rt << 1 | 1].lc) ans--;
    }
    //push_up(rt);
    return ans;
}
void change(int u, int v, int x)
{
    int tu = top[u], tv = top[v];
    while(tu != tv)
    {
        if(dep[tu] < dep[tv])
        swap(tu, tv), swap(u, v);
        update(1, pos[tu], pos[u], x);
        u = fa[tu];
        tu = top[u];
    }
    if(u == v) return ;
    if(dep[u] > dep[v]) swap(u, v);
    update(1, pos[son[u]], pos[v], x);
}
int find(int u, int v)
{
    int ans = 0;
    int tu = top[u], tv = top[v];
    int pre_u = u, pre_v = v, puc, pvc, lc, rc;//pre_u_color
    puc = pvc = -1;
    while(tu != tv)
    {
        if(dep[tu] < dep[tv])
        {
            swap(tu, tv), swap(u, v);
            swap(pre_u, pre_v), swap(puc, pvc);
        }
        ans += query(1, pos[tu], pos[u], lc, rc);
        if(rc == puc) ans--;
        puc = lc;
        pre_u = tu;
        u = fa[tu];
        tu = top[u];
    }
    if(u == v)
    {
        if(puc != -1 && puc == pvc) ans--;
        return ans;
    }
    if(dep[u] > dep[v])
    {
        swap(u, v);
        swap(pre_u, pre_v), swap(puc, pvc);
    }
    ans += query(1, pos[son[u]], pos[v], lc, rc);
    if(lc == puc) ans--;
    if(rc == pvc) ans--;
    return ans;
}
int main()
{
    int n, m, u, v, w;
    char s[10];
    while(~scanf("%d %d", &n, &m))
    {
        init(n);
        for(int i = 1; i < n; i++)
        {
            scanf("%d %d %d", &u, &v, &w);
            mp[u].push_back(v);
            mp[v].push_back(u);
            g[i].u = u, g[i].v = v, g[i].w = w;
        }
        dep[1] = fa[1] = 0;
        dfs1(1); dfs2(1, 1);
        for(int i = 1; i < n; i++)
        {
            if(dep[g[i].u] > dep[g[i].v])
            swap(g[i].u, g[i].v);
            color[pos[g[i].v]] = g[i].w;
        }
        //cout << tid << endl;
        build(1, tid, 1);
        while(m--)
        {
            scanf(" %s", s);
            if(s[0] == 'Q')
            {
                scanf("%d %d", &u, &v);
                printf("%d\n", find(u, v));
            }
            else
            {
                scanf("%d %d %d", &u, &v, &w);
                change(u, v, w);
            }
        }
    }
     return 0;
}