洛谷P4242 树上的毒瘤

解：首先有个套路是一条边的权值是[两端点颜色不同]。这个用树剖直接维护，支持修改。

每次询问建虚树，查询虚树上每条边的权值。然后树形DP，用开店的方法，每个点链加链查。

 #include <bits/stdc++.h>

 #define forson(x, i) for(int i = e[x]; i; i = edge[i].nex)

 typedef long long LL;
 const int N = ;

 struct Edge {
     int nex, v;
     LL len;
 }edge[N << ], EDGE[N]; int tp, TP;

 int e[N], top[N], fa[N], son[N], siz[N], d[N], pos[N], id[N], num, val[N], n, imp2[N];
 int sum[N << ], lc[N << ], rc[N << ], tag[N << ];
 int imp[N], K, stk[N], Top, RT, Time, E[N], vis[N], use[N], DEEP[N];
 LL SIZ[N], ans[N], D[N];

 inline void add(int x, int y) {
     tp++;
     edge[tp].v = y;
     edge[tp].nex = e[x];
     e[x] = tp;
     return;
 }

 /// ------------------- tree 1 -------------------------

 void DFS_1(int x, int f) { /// get fa son siz d
     fa[x] = f;
     siz[x] = ;
     d[x] = d[f] + ;
     forson(x, i) {
         int y = edge[i].v;
         if(y == f) continue;
         DFS_1(y, x);
         siz[x] += siz[y];
         if(siz[y] > siz[son[x]]) {
             son[x] = y;
         }
     }
     return;
 }

 void DFS_2(int x, int f) { /// get top pos id
     top[x] = f;
     pos[x] = ++num;
     id[num] = x;
     if(son[x]) DFS_2(son[x], f);
     forson(x, i) {
         int y = edge[i].v;
         if(y == fa[x] || y == son[x]) continue;
         DFS_2(y, y);
     }
     return;
 }

 /// ------------------ seg 1 ----------------------

 #define ls (o << 1)
 #define rs (o << 1 | 1)

 inline void pushup(int o) {
     lc[o] = lc[ls];
     rc[o] = rc[rs];
     sum[o] = sum[ls] + sum[rs] + (rc[ls] != lc[rs]);
     return;
 }

 inline void pushdown(int o) {
     if(tag[o] != -) {
         lc[ls] = rc[ls] = tag[ls] = tag[o];
         lc[rs] = rc[rs] = tag[rs] = tag[o];
         sum[ls] = sum[rs] = ;
         tag[o] = -;
     }
     return;
 }

 #undef ls
 #undef rs

 void build(int l, int r, int o) {
     if(l == r) {
         lc[o] = rc[o] = val[id[r]];
         sum[o] = ;
         return;
     }
     int mid = (l + r) >> ;
     build(l, mid, o << );
     build(mid + , r, o <<  | );
     pushup(o);
     return;
 }

 void change(int L, int R, int v, int l, int r, int o) {
     if(L <= l && r <= R) {
         lc[o] = rc[o] = tag[o] = v;
         sum[o] = ;
         return;
     }
     int mid = (l + r) >> ;
     pushdown(o);
     if(L <= mid) change(L, R, v, l, mid, o << );
     if(mid < R) change(L, R, v, mid + , r, o <<  | );
     pushup(o);
     return;
 }

 int ask(int p, int l, int r, int o) {
     if(l == r) return lc[o];
     int mid = (l + r) >> ;
     pushdown(o);
     if(p <= mid) return ask(p, l, mid, o << );
     else return ask(p, mid + , r, o <<  | );
 }

 int getSum(int L, int R, int l, int r, int o) {
     if(L <= l && r <= R) {
         return sum[o];
     }
     pushdown(o);
     int mid = (l + r) >> ;
     if(R <= mid) return getSum(L, R, l, mid, o << );
     if(mid < L) return getSum(L, R, mid + , r, o <<  | );
     return getSum(L, R, l, mid, o << ) + getSum(L, R, mid + , r, o <<  | ) + (rc[o << ] != lc[o <<  | ]);
 }

 inline int lca(int x, int y) {
     while(top[x] != top[y]) {
         if(d[top[x]] < d[top[y]])
             y = fa[top[y]];
         else
             x = fa[top[x]];
     }
     return d[x] < d[y] ? x : y;
 }

 inline int getLen(int x, int z) {
     //printf("getLen %d %d \n", x, z);
     int col = ask(pos[x], , n, ), ans = ;
     while(top[x] != top[z]) {
         ans += (col != ask(pos[x], , n, ));
         ans += getSum(pos[top[x]], pos[x], , n, );
         //printf("x = %d top[x] = %d col = %d  ans = %d \n", x, top[x], col, ans);
         col = ask(pos[top[x]], , n, );
         x = fa[top[x]];
     }
     ans += (col != ask(pos[x], , n, ));
     //printf("%d != %d \n", col, ask(pos[x], 1, n, 1));
     ans += getSum(pos[z], pos[x], , n, );
     //printf("return ans = %d \n", ans);
     return ans;
 }

 inline void Change(int x, int y, int v) {
     while(top[x] != top[y]) {
         if(d[top[x]] > d[top[y]]) {
             change(pos[top[x]], pos[x], v, , n, );
             x = fa[top[x]];
         }
         else {
             change(pos[top[y]], pos[y], v, , n, );
             y = fa[top[y]];
         }
     }
     if(d[x] < d[y]) std::swap(x, y);
     change(pos[y], pos[x], v, , n, );
     return;
 }

 /// ------------------- tree 2 ----------------------

 inline void ADD(int x, int y) {
     TP++;
     EDGE[TP].v = y;
     EDGE[TP].len = getLen(y, x);
     //printf("getLen %d %d = %d \n", y, x, EDGE[TP].len);
     EDGE[TP].nex = E[x];
     E[x] = TP;
     return;
 }

 inline bool cmp(const int &a, const int &b) {
     return pos[a] < pos[b];
 }

 inline void work(int x) {
     if(vis[x] == Time) return;
     vis[x] = Time;
     D[x] = E[x] = ;
     return;
 }

 inline void build_t() {
     TP = ;
     memcpy(imp + , imp2 + , K * sizeof(int));
     std::sort(imp + , imp + K + , cmp);
     stk[Top = ] = imp[];
     work(imp[]);
     for(int i = ; i <= K; i++) {
         int x = imp[i], y = lca(x, stk[Top]);
         work(x); work(y);
         while(Top >  && d[y] <= d[stk[Top - ]]) {
             ADD(stk[Top - ], stk[Top]);
             Top--;
         }
         if(y != stk[Top]) {
             ADD(y, stk[Top]);
             stk[Top] = y;
         }
         stk[++Top] = x;
     }
     while(Top > ) {
         ADD(stk[Top - ], stk[Top]);
         Top--;
     }
     RT = stk[Top];
     return;
 }

 void dfs_1(int x) { /// DP 1
     SIZ[x] = (use[x] == Time);
     for(int i = E[x]; i; i = EDGE[i].nex) {
         int y = EDGE[i].v;
         dfs_1(y);
         SIZ[x] += SIZ[y];
     }
     return;
 }

 void dfs_2(int x) { /// DP 2
     if(use[x] == Time) {
         ans[x] = D[x];
     }
     for(int i = E[x]; i; i = EDGE[i].nex) {
         int y = EDGE[i].v;
         D[y] = D[x] + SIZ[y] * EDGE[i].len;
         DEEP[y] = DEEP[x] + EDGE[i].len;
         //printf("dfs_2 D %d = %lld * %lld = %lld \n", y, SIZ[y], EDGE[i].len, D[y]);
         dfs_2(y);
     }
     return;
 }

 inline void cal() {
     build_t();
     dfs_1(RT);
     DEEP[RT] = ;
     dfs_2(RT);
     return;
 }

 int main() {
     memset(tag, -, sizeof(tag));
     int q;
     scanf("%d%d", &n, &q);
     for(int i = ; i <= n; i++) {
         scanf("%d", &val[i]);
     }
     for(int i = , x, y; i < n; i++) {
         scanf("%d%d", &x, &y);
         add(x, y); add(y, x);
     }
     DFS_1(, );
     DFS_2(, );
     build(, n, );

     for(int i = , f, x, y, z; i <= q; i++) {
         scanf("%d%d", &f, &x);
         if(f == ) {
             scanf("%d%d", &y, &z);
             Change(x, y, z);
         }
         else {
             Time++;
             K = x;
             for(int j = ; j <= K; j++) {
                 scanf("%d", &imp2[j]);
                 use[imp2[j]] = Time;
             }
             cal();
             LL SUM = ;
             for(int i = ; i <= K; i++) {
                 SUM += DEEP[imp2[i]];
                 //printf("D %d = %lld \n", imp2[i], D[imp2[i]]);
             }
             //printf("SUM = %lld \n", SUM);
             for(int i = ; i <= K; i++) {
                 printf("%lld ", SUM + K * DEEP[imp2[i]] -  * ans[imp2[i]] + K);
             }
             puts("");
         }
     }
     return ;
 }

AC代码