Rt, 调了一个下午了,救救孩子吧。
#include <bits/stdc++.h>
// #define int long long
using namespace std;
int N, M;
int w[100005], neww[100005];
struct Edge {
int nxt, to;
} edge[200005];
int cnt, head[100005];
void add(int u, int v)
{
++ cnt;
edge[cnt].nxt = head[u];
edge[cnt].to = v;
head[u] = cnt;
}
int father[100005], depth[100005], heavyson[100005], sz[100005];
void dfs1(int u, int fa)
{
father[u] = fa;
depth[u] = depth[fa] + 1;
sz[u] = 1;
for(int i = head[u]; i; i = edge[i].nxt)
{
int v = edge[i].to;
if(v == fa) continue;
dfs1(v, u);
sz[u] += sz[v];
if(sz[heavyson[u]] < sz[v]) heavyson[u] = v;
}
}
int dfncnt, dfn[100005], linktop[100005];
void dfs2(int u, int lktp)
{
dfn[u] = ++ dfncnt;
neww[dfn[u]] = w[u];
linktop[u] = lktp;
if(!heavyson[u]) return;
dfs2(heavyson[u], lktp);
for(int i = head[u]; i; i = edge[i].nxt)
{
int v = edge[i].to;
if(v == father[u] || v == heavyson[u]) continue;
dfs2(v, v);
}
}
#define lson (p << 1)
#define rson ((p << 1) | 1)
struct SegmentTree {
int l, r, left_end_color, right_end_color, segmentcnt, let_lazytag;
public:
SegmentTree()
{
this->l = 0ll, this->r = 0ll, this->left_end_color = 0ll, this->right_end_color = 0ll, this->segmentcnt = 0ll, let_lazytag = 0ll;
}
} tree[100000 << 2 + 5];
SegmentTree lquery, rquery;
SegmentTree merge_query(SegmentTree a, SegmentTree b)
{
if(!a.segmentcnt) return b;
if(!b.segmentcnt) return a;
SegmentTree res;
res.left_end_color = a.left_end_color, res.right_end_color = b.right_end_color;
res.segmentcnt = a.segmentcnt + b.segmentcnt;
if(a.right_end_color == b.left_end_color) -- res.segmentcnt;
return res;
}
void merge(int p)
{
tree[p].left_end_color = tree[lson].left_end_color, tree[p].right_end_color = tree[rson].right_end_color;
tree[p].segmentcnt = tree[lson].segmentcnt + tree[rson].segmentcnt;
if(tree[lson].right_end_color == tree[rson].left_end_color) -- tree[p].segmentcnt;
}
void build(int l, int r, int p)
{
tree[p].l = l, tree[p].r = r;
if(l == r)
{
tree[p].segmentcnt = 1;
tree[p].left_end_color = tree[p].right_end_color = neww[l];
return;
}
int mid = (l + r) >> 1;
build(l, mid, lson);
build(mid + 1, r, rson);
merge(p);
}
void pushdown(int p)
{
tree[lson].segmentcnt = tree[rson].segmentcnt = 1;
tree[lson].left_end_color = tree[lson].right_end_color = tree[rson].left_end_color = tree[rson].right_end_color = tree[p].let_lazytag;
tree[p].let_lazytag = 0;
}
void modify_let(int lll, int rrr, int p, int val)
{
if(lll <= tree[p].l && rrr >= tree[p].r)
{
tree[p].segmentcnt = 1;
tree[p].let_lazytag = tree[p].left_end_color = tree[p].right_end_color = val;
return;
}
if(tree[p].let_lazytag) pushdown(p);
int mid = (tree[p].l + tree[p].r) >> 1;
if(lll <= mid) modify_let(lll, rrr, lson, val);
if(rrr > mid) modify_let(lll, rrr, rson, val);
merge(p);
}
void modify_let_simplepath(int x, int y, int val)
{
while(linktop[x] != linktop[y])
{
if(depth[linktop[x]] < depth[linktop[y]]) swap(x, y);
modify_let(dfn[linktop[x]], dfn[x], 1, val);
x = father[linktop[x]];
}
if(depth[x] > depth[y]) swap(x, y);
modify_let(dfn[x], dfn[y], 1, val);
}
SegmentTree query_segmentcnt(int lll, int rrr, int p)
{
if(tree[p].l >= lll && tree[p].r <= rrr) return tree[p];
if(tree[p].let_lazytag) pushdown(p);
int mid = (tree[p].l + tree[p].r) >> 1;
if(rrr <= mid) return query_segmentcnt(lll, rrr, lson);
if(lll > mid) return query_segmentcnt(lll, rrr, rson);
return merge_query(query_segmentcnt(lll, rrr, lson), query_segmentcnt(lll, rrr, rson));
}
int query_segmentcnt_simplepath(int x, int y)
{
lquery = SegmentTree(), rquery = SegmentTree();
while(linktop[x] != linktop[y])
{
if(depth[linktop[x]] < depth[linktop[y]])
{
rquery = merge_query(query_segmentcnt(dfn[linktop[y]], dfn[y], 1), rquery);
y = father[linktop[y]];
}
else
{
lquery = merge_query(query_segmentcnt(dfn[linktop[x]], dfn[x], 1), lquery);
x = father[linktop[x]];
}
}
if(depth[x] < depth[y]) rquery = merge_query(query_segmentcnt(dfn[x], dfn[y], 1), rquery);
else lquery = merge_query(query_segmentcnt(dfn[y], dfn[x], 1), lquery);
swap(lquery.left_end_color, lquery.right_end_color);
return merge_query(lquery, rquery).segmentcnt;
}
signed main()
{
scanf("%d %d", &N, &M);
for(int i = 1; i <= N; ++ i)
scanf("%d", &w[i]);
for(int i = 1; i < N; ++ i)
{
int u, v;
scanf("%d %d", &u, &v);
add(u, v);
add(v, u);
}
dfs1(1, 0);
dfs2(1, 1);
build(1, N, 1);
for(int i = 1; i <= M; ++ i)
{
char op;
cin >> op;
if(op == 'C')
{
int a, b, c;
scanf("%d %d %d", &a, &b, &c);
modify_let_simplepath(a, b, c);
}
if(op == 'Q')
{
int a, b;
scanf("%d %d", &a, &b);
printf("%d\n", query_segmentcnt_simplepath(a, b));
}
}
return 0;
}