HYSBZ 2243 染色 (树链剖分)
HYSBZ 2243 染色
树链剖分,关键在于线段树的维护,对于每一个结点要记录下最左边和最右边的颜色。合并的时候。假设颜色同样那么颜色段要减1
代码:
#include <cstdio> #include <cstring> #include <vector> #include <algorithm> using namespace std; const int N = 100005; int dep[N], fa[N], son[N], sz[N], top[N], id[N], idx, val[N], val2[N]; int first[N], next[N * 2], vv[N * 2], en; void init() { en = 0; idx = 0; memset(first, -1, sizeof(first)); } void add_Edge(int u, int v) { vv[en] = v; next[en] = first[u]; first[u] = en++; } void dfs1(int u, int f, int d) { dep[u] = d; sz[u] = 1; fa[u] = f; son[u] = 0; for (int i = first[u]; i + 1; i = next[i]) { int v = vv[i]; if (v == f) continue; dfs1(v, u, d + 1); sz[u] += sz[v]; if (sz[son[u]] < sz[v]) son[u] = v; } } void dfs2(int u, int tp) { id[u] = ++idx; top[u] = tp; if (son[u]) dfs2(son[u], tp); for (int i = first[u]; i + 1; i = next[i]) { int v = vv[i]; if (v == fa[u] || v == son[u]) continue; dfs2(v, v); } } #define lson(x) ((x<<1)+1) #define rson(x) ((x<<1)+2) struct Node { int l, r, lc, rc, cnt, setv; Node() {cnt = 0;} } node[N * 4]; void pushup(int x) { node[x].cnt = node[lson(x)].cnt + node[rson(x)].cnt; if (node[lson(x)].rc == node[rson(x)].lc) node[x].cnt--; node[x].lc = node[lson(x)].lc; node[x].rc = node[rson(x)].rc; } void pushdown(int x) { if (node[x].setv == -1) return; node[lson(x)].setv = node[lson(x)].lc = node[lson(x)].rc = node[x].setv; node[rson(x)].setv = node[rson(x)].lc = node[rson(x)].rc = node[x].setv; node[lson(x)].cnt = node[rson(x)].cnt = 1; node[x].setv = -1; } void build(int l, int r, int x = 0) { node[x].l = l; node[x].r = r; node[x].setv = -1; if (l == r) { node[x].lc = node[x].rc = val[l]; node[x].cnt = 1; return; } int mid = (l + r) / 2; build(l, mid, lson(x)); build(mid + 1, r, rson(x)); pushup(x); } void add(int l, int r, int v, int x = 0) { if (node[x].l >= l && node[x].r <= r) { node[x].cnt = 1; node[x].setv = node[x].lc = node[x].rc = v; return; } pushdown(x); int mid = (node[x].l + node[x].r) / 2; if (l <= mid) add(l, r, v, lson(x)); if (r > mid) add(l, r, v, rson(x)); pushup(x); } Node merge(Node a, Node b) { Node ans; if (a.cnt == 0) return b; if (b.cnt == 0) return a; ans.lc = a.lc; ans.rc = b.rc; ans.cnt = a.cnt + b.cnt; if (a.rc == b.lc) ans.cnt--; return ans; } Node query(int l, int r, int x = 0) { if (node[x].l >= l && node[x].r <= r) { return node[x]; } pushdown(x); int mid = (node[x].l + node[x].r) / 2; Node ans; if (l <= mid) ans = merge(query(l, r, lson(x)), ans); if (r > mid) ans = merge(ans, query(l, r, rson(x))); pushup(x); return ans; } void gao1(int u, int v, int val) { int tp1 = top[u], tp2 = top[v]; while (tp1 != tp2) { if (dep[tp1] < dep[tp2]) { swap(tp1, tp2); swap(u, v); } add(id[tp1], id[u], val); u = fa[tp1]; tp1 = top[u]; } if (dep[u] > dep[v]) swap(u, v); add(id[u], id[v], val); } int gao2(int u, int v) { int tp1 = top[u], tp2 = top[v]; Node ans1, ans2; while (tp1 != tp2) { if (dep[tp1] < dep[tp2]) { swap(tp1, tp2); swap(u, v); swap(ans1, ans2); } ans1 = merge(query(id[tp1], id[u]), ans1); u = fa[tp1]; tp1 = top[u]; } if (dep[u] > dep[v]) { swap(ans1, ans2); swap(u, v); } ans2 = merge(query(id[u], id[v]), ans2); if (ans1.cnt == 0) return ans2.cnt; if (ans2.cnt == 0) return ans1.cnt; int ans = ans1.cnt + ans2.cnt; if (ans1.lc == ans2.lc) ans--; return ans; } int n, m; int main() { while (~scanf("%d%d", &n, &m)) { init(); for (int i = 1; i <= n; i++) scanf("%d", &val2[i]); int u, v; for (int i = 1; i < n; i++) { scanf("%d%d", &u, &v); add_Edge(u, v); add_Edge(v, u); } dfs1(1, 0, 1); dfs2(1, 1); for (int i = 1; i <= n; i++) val[id[i]] = val2[i]; build(1, idx); char Q[10]; int a, b, c; while (m--) { scanf("%s%d%d", Q, &a, &b); if (Q[0] == 'Q') printf("%d\n", gao2(a, b)); else { scanf("%d", &c); gao1(a, b, c); } } } return 0; }