Codeforces 633G Yash And Trees bitset + 线段树

Yash And Trees

用bitset维护每个节点拥有哪些数。

#include<bits/stdc++.h>
#define LL long long
#define LD long double
#define ull unsigned long long
#define fi first
#define se second
#define mk make_pair
#define PLL pair<LL, LL>
#define PLI pair<LL, int>
#define PII pair<int, int>
#define SZ(x) ((int)x.size())
#define ALL(x) (x).begin(), (x).end()
#define fio ios::sync_with_stdio(false); cin.tie(0);

using namespace std;

const int N = 1e5 + 7;
const int inf = 0x3f3f3f3f;
const LL INF = 0x3f3f3f3f3f3f3f3f;
const int mod = 1e9 + 7;
const double eps = 1e-8;
const double PI = acos(-1);

template<class T, class S> inline void add(T& a, S b) {a += b; if(a >= mod) a -= mod;}
template<class T, class S> inline void sub(T& a, S b) {a -= b; if(a < 0) a += mod;}
template<class T, class S> inline bool chkmax(T& a, S b) {return a < b ? a = b, true : false;}
template<class T, class S> inline bool chkmin(T& a, S b) {return a > b ? a = b, true : false;}

int n, m, q, a[N];
int in[N], ot[N], b[N], idx;
bitset<1000> tmp[2];
bitset<1000> prime;
bitset<1000> ans;

vector<int> G[N];

#define lson l, mid, rt << 1
#define rson mid + 1, r, rt << 1 | 1
struct setmentTree {
    bitset<1000> a[N << 2];
    int lazy[N << 2];

    inline void gao(int rt, int c) {
        tmp[0] = (a[rt] << (1000 - m)) >> (1000 - c);
        tmp[1] = (a[rt] << (c + 1000 - m)) >> (1000 - m);
        a[rt] = tmp[0] | tmp[1];
        lazy[rt] += c; if(lazy[rt] >= m) lazy[rt] -= m;
    }

    inline void push(int rt) {
        if(lazy[rt]) {
            gao(rt << 1, lazy[rt]);
            gao(rt << 1 | 1, lazy[rt]);
            lazy[rt] = 0;
        }
    }

    void build(int *b, int l, int r, int rt) {
        if(l == r) {
            a[rt][b[l]] = 1;
            return;
        }
        int mid = l + r >> 1;
        build(b, lson); build(b, rson);
        a[rt] = a[rt << 1] | a[rt << 1 | 1];
    }
    
    void update(int L, int R, int val, int l, int r, int rt){
        if(R < l || r < L || R < L) return;
        if(L <= l && r <= R) {
            gao(rt, val);
            return;
        }
        push(rt);
        int mid = l + r >> 1;
        update(L, R, val, lson);
        update(L, R, val, rson);
        a[rt] = a[rt << 1] | a[rt << 1 | 1];
    }
    bitset<1000> query(int L, int R, int l, int r, int rt) {
        if(L <= l && r <= R) return a[rt];
        push(rt);
        int mid = l + r >> 1;
        bitset<1000> ans;
        if(L <= mid) ans |= query(L, R, lson);
        if(R > mid) ans |= query(L, R, rson);
        return ans;
    }
} Tree;

void dfs(int u, int fa) {
    in[u] = ++idx;
    b[idx] = a[u];
    for(auto &v : G[u])
        if(v != fa) dfs(v, u);
    ot[u] = idx;
}

bool isPrime(int x) {
    for(int i = 2; i * i <= x; i++)
        if(x % i == 0) return false;
    return true;
}

int main() {
    scanf("%d%d%", &n, &m);
    for(int i = 2; i < m; i++) prime[i] = isPrime(i);
    for(int i = 1; i <= n; i++) scanf("%d", &a[i]), a[i] %= m;
    for(int i = 1; i < n; i++) {
        int u, v;
        scanf("%d%d", &u, &v);
        G[u].push_back(v);
        G[v].push_back(u);
    }
    dfs(1, 0);
    Tree.build(b, 1, n, 1);
    scanf("%d", &q);
    while(q--) {
        int op, v, x;
        scanf("%d", &op);
        if(op == 1) {
            scanf("%d%d", &v, &x);
            x %= m;
            Tree.update(in[v], ot[v], x, 1, n, 1);
        } else {
            scanf("%d", &v);
            ans = Tree.query(in[v], ot[v], 1, n, 1);
            printf("%d\n", (ans & prime).count());
        }
    }
    return 0;
}

/*
*/