Codeforces 633G 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 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,rson);
        a[rt] = a[rt << 1] | a[rt << 1 | 1];
    }
    bitset<1000> query(int L,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,lson);
        if(R > mid) ans |= query(L,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,1);
        } else {
            scanf("%d",&v);
            ans = Tree.query(in[v],1);
            printf("%dn",(ans & prime).count());
        }
    }
    return 0;
}

/*
*/