[Bzoj2243][SDOI2011]染色(线段树树剖)

题目链接：https://www.lydsy.com/JudgeOnline/problem.php?id=2243

线段树+树链剖分，在线段树需要每次用lt和rt两个数组记录当前区间的左右边界的颜色，向上更新时需要判断左区间的右边界是否和右区间的左边界相等。在剖分求LCA的过程中需要在求值之后查询与下一次求值的边界是否相等。

  1 #include<bits/stdc++.h>
  2 #define lson l,mid,i<<1
  3 #define rson mid+1,r,i<<1|1
  4 using namespace std;
  5 typedef long long ll;
  6 const int maxn = 200005;
  7 const int INF = 2e9;
  8 struct node {
  9     int s, e, next;
 10 }edge[maxn * 2];
 11 int n, m;
 12 int son[maxn], top[maxn], tid[maxn], fat[maxn], siz[maxn], dep[maxn], rak[maxn];
 13 int head[maxn], len, dfx;
 14 //siz保存以i为根的子树节点个数，top保存i节点所在链的顶端节点,son保存i节点的重儿子,fat保存i节点的父亲节点
 15 //dep保存i节点的深度(根为1),,tid保存i节点dfs后的新编号，rak保存新编号i对应的节点(rak[i]=j,tid[j]=i)。
 16 void init() {
 17     memset(head, -1, sizeof(head));
 18     len = 0, dfx = 0;
 19 }
 20 void add(int s, int e) {//邻接表存值
 21     edge[len].s = s;
 22     edge[len].e = e;
 23     edge[len].next = head[s];
 24     head[s] = len++;
 25 }
 26 //搜出每个节点的siz,son,fat,dep
 27 void dfs1(int x, int fa, int d) {
 28     siz[x] = 1, son[x] = -1, fat[x] = fa, dep[x] = d;
 29     for (int i = head[x]; i != -1; i = edge[i].next) {
 30         int y = edge[i].e;
 31         if (y == fa)
 32             continue;
 33         dfs1(y, x, d + 1);
 34         siz[x] += siz[y];
 35         if (son[x] == -1 || siz[y] > siz[son[x]])
 36             son[x] = y;
 37     }
 38 }
 39 //搜出每个节点的top,tid,rak
 40 void dfs2(int x, int c) {
 41     top[x] = c;
 42     tid[x] = ++dfx;
 43     rak[dfx] = x;
 44     if (son[x] == -1)
 45         return;
 46     dfs2(son[x], c);
 47     for (int i = head[x]; i != -1; i = edge[i].next) {
 48         int y = edge[i].e;
 49         if (y == fat[x] || y == son[x])
 50             continue;
 51         dfs2(y, y);
 52     }
 53 }
 54 int a[maxn];
 55 int cr[maxn * 2];
 56 int rt[maxn * 2];
 57 int lt[maxn * 2];
 58 int lazy[maxn * 2];
 59 void up(int i) {
 60     lt[i] = lt[i << 1], rt[i] = rt[i << 1 | 1];
 61     cr[i] = cr[i << 1] + cr[i << 1 | 1];
 62     if (lt[i << 1 | 1] == rt[i << 1])
 63         cr[i]--;
 64 }
 65 void down(int i) {
 66     if (lazy[i] != -1) {
 67         lt[i << 1] = lt[i << 1 | 1] = rt[i << 1] = rt[i << 1 | 1] = lazy[i];
 68         cr[i << 1] = cr[i << 1 | 1] = cr[i];
 69         lazy[i << 1] = lazy[i << 1 | 1] = lazy[i];
 70         lazy[i] = -1;
 71     }
 72 }
 73 void build(int l, int r, int i) {
 74     lazy[i] = -1;
 75     if (l == r) {
 76         cr[i] = 1;
 77         rt[i] = a[rak[l]];
 78         lt[i] = a[rak[l]];
 79         return;
 80     }
 81     int mid = (l + r) >> 1;
 82     build(lson);
 83     build(rson);
 84     up(i);
 85 }
 86 void update(int L, int R, int k, int l, int r, int i) {
 87     if (L <= l && r <= R) {
 88         cr[i] = 1;
 89         lazy[i] = k;
 90         lt[i] = rt[i] = k;
 91         return;
 92     }
 93     down(i);
 94     int mid = (l + r) >> 1;
 95     if (L <= mid)
 96         update(L, R, k, lson);
 97     if (R > mid)
 98         update(L, R, k, rson);
 99     up(i);
100 }
101 int qquery(int L, int R, int l, int r, int i) {
102     if (L <= l && r <= R) return cr[i];
103     down(i);
104     int mid = (l + r) >> 1;
105     if (R <= mid) return qquery(L, R, lson);
106     if (L > mid) return qquery(L, R, rson);
107     int ans1 = qquery(L, R, lson);
108     int ans2 = qquery(L, R, rson);
109     int ans = ans1 + ans2;
110     if (rt[i << 1] == lt[i << 1 | 1]) ans--;
111     return ans;
112 }
113 
114 int dquery(int k, int l, int r, int i) {
115     if (l == r) {
116         return lt[i];
117     }
118     int mid = (l + r) / 2;
119     down(i);
120     if (k <= mid)
121         return dquery(k, lson);
122     else
123         return dquery(k, rson);
124 }
125 int solve(int x, int y, int w, int flg) {
126     int ans = 0;
127     while (top[x] != top[y]) {
128         if (dep[top[x]] < dep[top[y]])
129             swap(x, y);
130         if (!flg)
131             update(tid[top[x]], tid[x], w, 1, n, 1);
132         else {
133             ans += qquery(tid[top[x]], tid[x], 1, n, 1);
134             if (dquery(tid[top[x]], 1, n, 1) == dquery(tid[fat[top[x]]], 1, n, 1))
135                 ans--;
136         }
137         x = fat[top[x]];
138     }
139     if (dep[x] < dep[y])
140         swap(x, y);
141     if (!flg)
142         update(tid[y], tid[x], w, 1, n, 1);
143     else {
144         ans += qquery(tid[y], tid[x], 1, n, 1);
145         return ans;
146     }
147 }
148 int main() {
149     while (scanf("%d%d", &n, &m) != EOF) {
150         for (int i = 1; i <= n; i++)
151             scanf("%d", &a[i]);
152         init();
153         int x, y, z;
154         for (int i = 0; i < n - 1; i++) {
155             scanf("%d%d", &x, &y);
156             add(x, y);
157             add(y, x);
158         }
159         dfs1(1, 0, 1);
160         dfs2(1, 1);
161         build(1, n, 1);
162         while (m--) {
163             char s[10];
164             scanf("%s", s);
165             if (s[0] == 'C') {
166                 scanf("%d%d%d", &x, &y, &z);
167                 solve(x, y, z, 0);
168             }
169             else {
170                 scanf("%d%d", &x, &y);
171                 printf("%d\n", solve(x, y, 0, 1));
172             }
173         }
174     }
175 }