「日常训练」Kefa and Park（Codeforces Round #321 Div. 2 C）

题意与分析（CodeForces 580C）

代码

#include <bits/stdc++.h>
#define MP make_pair
#define PB emplace_back
#define fi first
#define se second
#define ZERO(x) memset((x), 0, sizeof(x))
#define ALL(x) (x).begin(),(x).end()
#define rep(i, a, b) for (repType i = (a); i <= (b); ++i)
#define per(i, a, b) for (repType i = (a); i >= (b); --i)
#define QUICKIO                  \
    ios::sync_with_stdio(false); \
    cin.tie(0);                  \
    cout.tie(0);
using namespace std;
using ll=long long;
using repType=ll;
const int MAXN=100000+5;

bool red[MAXN];
bool ok[MAXN];
bool vis[MAXN];
vector<int> G[MAXN],nG[MAXN];
int n;

int dfs(int x)
{
    if(nG[x].size()==0)
    {
        //cout<<x<<endl;
        return ok[x];
    }
    else
    {
        int ans=0;
        rep(i,0,int(nG[x].size())-1)
            ans+=dfs(nG[x][i]);
        return ans;
    }
}

int bfs(int m)
{
    ZERO(ok);
    queue<pair<int,int> > q;
    q.push(MP(1,0));

    vis[1]=true;
    while(!q.empty())
    {
        auto now=q.front(); q.pop();
        //cout<<now.fi<<" "<<now.se<<" "<<red[now.fi]<<endl;
        if(now.se+red[now.fi]>m) continue;
        else ok[now.fi]=true;
        rep(i,0,int(G[now.fi].size())-1)
        {
            int v=G[now.fi][i];
            if(!vis[v])
            {
                vis[v]=true;
                nG[now.fi].PB(v);
                if(red[now.fi])
                    q.push(MP(v,now.se+1));
                else q.push(MP(v,0));
            }
        }
    }
    return dfs(1);
}

int main()
{
    int m;
    cin>>n>>m;
    rep(i,1,n) cin>>red[i];
    rep(i,1,n)
    {
        int x,y; cin>>x>>y;
        G[x].PB(y);
        G[y].PB(x);
    }
    ZERO(vis);
    cout<<bfs(m)<<endl;
       
    return 0;
}