[HDU 5893] List wants to travel （树链剖分+区间合并）

链接

HDU 5893

---

题意

沈阳网赛的第二题。。。
给出一棵N个点的树，边含权，然后是M个操作，每个操作“Change a b c”改变节点a到节点b的路径上权值为c，“Query a b”询问节点a到节点b路径上有多少个权值相同的段。
差不多是裸的树链剖分，维护好树链以后就可以利用线段树记录区间左右端点值做合并了。
这里需要注意的是树上的链的合并计数也需要链的端点值，u、v两点可能需要跨很多链最终才到一条重链上，所以在跨链的时候需要判断新链的底端和旧链的顶端是否是同一种颜色，是则-1，并且不断更新链的端点。
这里总结一下线段树的lazy相关的操作：
成段更新（区间染色）这个操作，需要用lazy记录，这样在push_down的时候需要将标记往下推，并且更新推到的区间的信息，如果最终需要更改某个区间，也就是update操作，则需要push_up操作做合并，因为更改的下端值会影响到路径到更改点的所有区间，这对于query操作往往是不需要的，因为没有区间改变就不会影响到lazy层上侧的区间。

---

代码

#include <cstdio>
#include <iostream>
#include <algorithm>
#include <vector>
using namespace std;
#define next nnext
#define maxn (40040)
struct edge { int u, v, w; } e[maxn];
vector<int> son[maxn];
int fa[maxn], sz[maxn], dep[maxn], next[maxn];
void dfs1(int u, int f, int d)
{
    fa[u] = f, sz[u] = 1, dep[u] = d, next[u] = 0;
    for(int i = 0, v; i < son[u].size(); i++)
    if((v = son[u][i]) != f) {
        dfs1(v, u, d + 1);
        sz[u] += sz[v];
        if(sz[v] > sz[next[u]]) next[u] = v;
    }
}
int top[maxn], idx[maxn], w[maxn], cnt;
void dfs2(int u, int tp)
{
    top[u] = tp;
    idx[u] = ++cnt;
    if(!next[u]) return ;
    dfs2(next[u], tp);
    for(int i = 0, v; i < son[u].size(); i++)
    if((v = son[u][i]) != fa[u] && v != next[u]) {
        dfs2(v, v);
    }
}
#define lson l, m, rt<<1
#define rson m+1, r, rt<<1|1
int sum[maxn << 2], le[maxn << 2], re[maxn << 2], lazy[maxn << 2];
inline void push_up(int rt)
{
    sum[rt] = sum[rt<<1] + sum[rt<<1|1] - (re[rt<<1] == le[rt<<1|1]);
    le[rt] = le[rt<<1], re[rt] = re[rt<<1|1];
}
inline void push_down(int rt)
{
    if(lazy[rt] == -1) return ;
    le[rt<<1] = re[rt<<1] = lazy[rt<<1] = lazy[rt];
    le[rt<<1|1] = re[rt<<1|1] = lazy[rt<<1|1] = lazy[rt];
    sum[rt<<1] = sum[rt<<1|1] = 1;
    lazy[rt] = -1;
}
void build(int l, int r, int rt)
{
    lazy[rt] = -1;
    if(l == r) { le[rt] = re[rt] = w[l]; sum[rt] = 1; return; }
    int m = (l + r) >> 1;
    build(lson), build(rson);
    push_up(rt);
}
void update(int L, int R, int val, int l, int r, int rt)
{
    if(L <= l && r <= R) { le[rt] = re[rt] = lazy[rt] = val; sum[rt] = 1; return; }
    push_down(rt);
    int m = (l + r) >> 1;
    if(L <= m) update(L, R, val, lson);
    if(R > m)  update(L, R, val, rson);
    push_up(rt);
}
int up, down;
int query(int L, int R, int l, int r, int rt)
{
    if(L <= l && r <= R) {
        if(L == l) up = le[rt];
        if(R == r) down = re[rt];
        return sum[rt];
    }
    push_down(rt);
    int m = (l + r) >> 1, ret = 0, fl = 0, fr = 0;
    if(L <= m) { ret += query(L, R, lson); fl = 1; }
    if(R > m)  { ret += query(L, R, rson); fr = 1; }
    return ret - (fl && fr && (re[rt<<1] == le[rt<<1|1]));
}
void change_tree(int u, int v, int val, int n)
{
    int topu = top[u], topv = top[v];
    while(topu != topv)
    {
        if(dep[topu] < dep[topv]) { swap(u, v); swap(topu, topv); }
        update(idx[topu], idx[u], val, 1, n, 1);
        u = fa[topu];
        topu = top[u];
    }
    if(u == v) return ;
    if(dep[u] > dep[v]) swap(u, v);
    update(idx[next[u]], idx[v], val, 1, n, 1);
}
int ask_tree(int u, int v, int n)
{
    int sumu = 0, sumv = 0, p, ret;
    int u_up = -1, v_up = -1;
    while(top[u] != top[v])
    {
        p = (dep[top[u]] > dep[top[v]]) ? u : v;
        ret = query(idx[top[p]], idx[p], 1, n, 1);
        if(p == u)
        {
            sumu += ret;
            sumu -= (u_up == down);
            u_up = up;
            u = fa[top[u]];
        }
        else
        {
            sumv += ret;
            sumv -= (v_up == down);
            v_up = up;
            v = fa[top[v]];
        }
    }
    if(u == v) return max(0, sumu + sumv - (u_up == v_up));
    if(dep[u] > dep[v])
    {
        ret = query(idx[next[v]], idx[u], 1, n, 1);
        return ret + sumv + sumu - (up == v_up) - (down == u_up);
    }
    else
    {

        ret = query(idx[next[u]], idx[v], 1, n, 1);
        return ret + sumv + sumu - (up == u_up) - (down == v_up);
    }
}
int main()
{
    int N, M;
    while(cin >> N >> M)
    {
        cnt = 0;
        for(int i = 1; i <= N; i++) son[i].clear();
        for(int i = 0; i < N-1; i++)
        {
            scanf("%d%d%d", &e[i].u, &e[i].v, &e[i].w);
            son[e[i].u].push_back(e[i].v);
            son[e[i].v].push_back(e[i].u);
        }
        dfs1(1, 0, 1);
        dfs2(1, 1);
        for(int i = 0; i < N-1; i++)
        {
            if(dep[e[i].u] < dep[e[i].v]) swap(e[i].u, e[i].v);
            w[idx[e[i].u]] = e[i].w;
        }
        build(1, N, 1);
        char op[10];
        int a, b, c;
        while(M--)
        {
            scanf("%s", op);
            if(op[0] == 'Q')
            {
                scanf("%d%d", &a, &b);
                printf("%d\n", ask_tree(a, b, N));
            }
            else
            {
                scanf("%d%d%d", &a, &b, &c);
                change_tree(a, b, c, N);
            }
        }
    }
    return 0;
}