染色

题目描述

思路

尽量不要直接使用getchar得到单个字符，容易出现runtime error， 可以使用scanf("%s", str), 直接判断str[0]是否是对应字符就可以了。

代码

#include <cstdio>
#include <cstring>
#define lc k<<1
#define rc k<<1|1

const int MAX = 100100;
int n, m, wt[MAX], ans;
int head[MAX], ver[MAX << 1], nt[MAX << 1], ht;
int fa[MAX], size[MAX], son[MAX], dep[MAX];
int top[MAX], dfn[MAX], tr[MAX], dt;
int lg[MAX << 2], rg[MAX << 2], sum[MAX << 2];
bool fg[MAX << 2];
char str[100];

int read() {
    int s = 0, f = 1;
    char ch = getchar();
    while (ch < '0' || ch > '9') {
        if (ch == '-') f = -1;
        ch = getchar();
    }
    while (ch >= '0' && ch <= '9') {
        s = s * 10 + ch - '0'; ch = getchar();
    }
    return s * f;
}
void write(int x) {
    if (x > 9) write(x / 10);
    putchar(x % 10 + '0');
}

void add(int x, int y) {
    nt[++ht] = head[x], head[x] = ht, ver[ht] = y;
}

void dfs1(int x, int u) {
    fa[x] = u;
    dep[x] = dep[u] + 1;
    size[x] = 1;
    for (int i = head[x], j; j = ver[i]; i = nt[i]) {
        if (j == u) continue;
        dfs1(j, x);
        size[x] += size[j];
        if (size[j] > size[son[x]]) son[x] = j;
    }
}

void dfs2(int x, int u) {
    top[x] = u;
    dfn[x] = ++dt;
    tr[dt] = x;
    if (son[x]) dfs2(son[x], u);
    for (int i = head[x], j; j = ver[i]; i = nt[i]) {
        if (!dfn[j]) dfs2(j, j);
    }
}
void pushup(int k, int l, int r) {
    lg[k] = lg[lc], rg[k] = rg[rc];
    sum[k] = sum[lc] + sum[rc];
    if (rg[lc] == lg[rc]) sum[k]--;  
}

void pushdown(int k, int l, int r, int mid) {
    if (!fg[k]) return;
    fg[lc] = fg[rc] = true;
    sum[lc] = sum[rc] = 1;
    lg[lc] = lg[rc] = rg[lc] = rg[rc] = lg[k];
    fg[k] = false;
}

void swap(int &x, int &y) {
    int t = x;
    x = y, y = t;
}

void build(int k, int l, int r) {
    if (l == r) {
        lg[k] = rg[k] = wt[tr[l]];
        sum[k] = 1;
        return;
    }
    int mid = l + r >> 1;
    build(lc, l, mid);
    build(rc, mid + 1, r);
    pushup(k, l, r);
}

int getColor(int k, int l, int r, int x) {
    if (l == r && x == l) return lg[k];
    int mid = l + r >> 1;
    pushdown(k, l, r, mid);
    if (x <= mid) return getColor(lc, l, mid, x);
    else return getColor(rc, mid + 1, r, x);
}

void change(int k, int l, int r, int x, int y, int z) {
    if (x <= l && r <= y) {
        fg[k] = true;
        sum[k] = 1;
        lg[k] = rg[k] = z;
        return;
    }
    int mid = l + r >> 1;
    pushdown(k, l, r, mid);
    if (x <= mid) change(lc, l, mid, x, y, z);
    if (y > mid) change(rc, mid + 1, r, x, y, z);
    pushup(k, l, r);
}

void query(int k, int l, int r, int x, int y) {
    if (x <= l && r <= y) {
        ans += sum[k];
        return;
    }
    int mid = l + r >> 1;
    pushdown(k, l, r, mid);
    if (x <= mid) query(lc, l, mid, x, y);
    if (y > mid) query(rc, mid + 1, r, x, y);
    if (x <= mid && y > mid && rg[lc] == lg[rc]) ans--;
}

void ask(int x, int y) {
    int fx = top[x], fy = top[y];
    while (fx != fy) {
        if (dep[fx] < dep[fy]) swap(x, y), swap(fx, fy);
        query(1, 1, n, dfn[fx], dfn[x]);
        if (getColor(1, 1, n, dfn[fx]) == getColor(1, 1, n, dfn[fa[fx]])) ans--;
        x = fa[fx], fx = top[x];
    }
    if (dep[x] > dep[y]) swap(x, y);
    query(1, 1, n, dfn[x], dfn[y]);
}

void askChange(int x, int y, int z) {
    int fx = top[x], fy = top[y];
    while (fx != fy) {
        if (dep[fx] < dep[fy]) swap(x, y), swap(fx, fy);
        change(1, 1, n, dfn[fx], dfn[x], z);
        x = fa[fx], fx = top[x];
    }
    if (dep[x] > dep[y]) swap(x, y);
    change(1, 1, n, dfn[x], dfn[y], z);
}

int main() {
    n = read(), m = read();
    for (int i = 1; i <= n; ++i) wt[i] = read();
    for (int i = 1, a, b; i < n; ++i) {
        a = read(), b = read();
        add(a, b), add(b, a);
    }
    dfs1(1, 0);
    dfs2(1, 1);
    build(1, 1, n);
    for (int i = 1, j, a, b, c; i <= m; ++i) {
        scanf("%s", str);
        if (str[0] == 'Q') {
            a = read(), b = read();
            ans = 0;
            ask(a, b);
            write(ans);
            puts(""); 
        } else {
            a = read(), b = read(), c = read();
            askChange(a, b, c);
        }
    }
    return 0;
}