Educational Codeforces Round 6 E. New Year Tree

問題

codeforces.com

解法

見た目が明らかに Euler Tour + 遅延評価セグ木 です。今回, 色の種類が 60 種類しかないので, セグメント木には, 2^c[i] (色 c[i] が含まれるかどうかのフラグ) というのの OR を取る形のものを保持しておけば良いです。

struct ST {
    vector<ll> seg, lazy;
    int size;

    ST(int n) {
        size = 1;
        while (size < n) size *= 2;
        seg.resize(size * 2);
        lazy.resize(size * 2, -1);
    }

    inline void push(int k, int l, int r) {
        if (lazy[k] != -1) {
            seg[k] = lazy[k];
            if (r - l > 1) {
                lazy[k * 2 + 1] = lazy[k];
                lazy[k * 2 + 2] = lazy[k];
            }
            lazy[k] = -1;
        }
    }

    inline ll merge(ll x, ll y) {
        return x | y;
    }

    void update(int a, int b, ll v, int k, int l, int r) {
        push(k, l, r);
        if (r <= a || b <= l) return;
        if (a <= l && r <= b) {
            lazy[k] = v;
            push(k, l, r);
        } else {
            update(a, b, v, k * 2 + 1, l, (l + r) / 2);
            update(a, b, v, k * 2 + 2, (l + r) / 2, r);
            seg[k] = merge(seg[k * 2 + 1], seg[k * 2 + 2]);
        }
    }

    void update(int a, int b, ll v) {
        return update(a, b, v, 0, 0, size);
    }

    ll query(int a, int b, int k, int l, int r) {
        push(k, l, r);
        if (r <= a || b <= l) return 0;
        if (a <= l && r <= b) return seg[k];
        ll x = query(a, b, k * 2 + 1, l, (l + r) / 2);
        ll y = query(a, b, k * 2 + 2, (l + r) / 2, r);
        return merge(x, y);
    }

    ll query(int a, int b) {
        return query(a, b, 0, 0, size);
    }
};

// Euler-Tour
const int MAXSIZE = 400020;
int BEGIN[MAXSIZE], END[MAXSIZE];
vector<int> euler_tour;
int K;
vector<vi> G;

void createEulerTour(int v, int p) {
    BEGIN[v] = K++;
    euler_tour.push_back(v);
    for (int el : G[v]) {
        if (el == p) continue;
        createEulerTour(el, v);
        euler_tour.push_back(v);
        K++;
    }
    END[v] = K;
}

int C[400400];

int main() {
    int n, m;
    scanf("%d %d", &n, &m);
    G.resize(n);
    for (int i = 0; i < n; i++) {
        scanf("%d", C+i);
        C[i]--;
    }
    for (int i = 0; i < n-1; i++) {
        int x, y;
        scanf("%d %d", &x, &y);
        x--; y--;
        G[x].push_back(y);
        G[y].push_back(x);
    }
    createEulerTour(0, -1);
    ST segs(euler_tour.size());
    for (int i = 0; i < n; i++) {
        segs.update(BEGIN[i], BEGIN[i]+1, 1ll<<C[i]);
    }
    while (m--) {
        int t;
        scanf("%d", &t);
        if (t==1) {
            int v, c;
            scanf("%d %d", &v, &c);
            v--; c--;
            segs.update(BEGIN[v], END[v], 1ll<<c);
        } else {
            int v;
            scanf("%d", &v);
            v--;
            int ans = __builtin_popcountll(segs.query(BEGIN[v], END[v]));
            printf("%d\n", ans);
        }
    }
    return 0;
}