Comet OJ - Contest #5 迫真小游戏

https://www.cometoj.com/contest/46/problem/C?problem_id=2028

思路：一层一层设置；优先设置编号小的，每次设置一个，将其限制的层减少，如果该层没限制了，就可以设置该层； 

#include <iostream>
#include <cstring>
#include <algorithm>
#include <cstdio>
#include <queue>
#include <map>
#include <set>
#include <stack>
using namespace std;

#define sfi(x) scanf("%d",&x)
#define sfc(x) scanf("%c",x)
#define sfl(x) scanf("%lld",&x)
#define sfs(x) scanf("%s",x)

#define rint register int
#define pb push_back
#define fl() printf("flag!\n")
#define INF 0x3f3f3f3f
#define ll long long
#define mem(x,y) memset(x,y,sizeof(x))
#define FAST_IO ios::sync_with_stdio(false);cin.tie(0);cout.tie(0)

const int maxn=5e5+9;

vector<int>g[maxn],vec[maxn];
int num[maxn],a[maxn];
int n;
int Maxd=0;
void dfs(int u,int fa,int d)
{
    vec[d].pb(u);
    Maxd=max(d,Maxd);
    for(int i=0;i<g[u].size();i++)
    {
        int v=g[u][i];
        if(v==fa) continue;
        dfs(v,u,d+1);
    }
}

int main()
{
    //FAST_IO;
    //freopen("input.txt","r",stdin);

    cin>>n;
    int u,v;
    for(int i=1;i<=n-1;i++)
    {
        cin>>u>>v;
        g[u].pb(v);
        g[v].pb(u);
    }
    dfs(1,0,1);
    for(int i=1;i<=n;i++)
    {
        int x;
        cin>>x;
        x++;
        num[x]++;
        a[i]=x;
    }

    priority_queue<int,vector<int>,greater<int> >q;
    q.push(1);
    int cnt=2;
    vector<int>ans;
    while(!q.empty())
    {
        int tmp=q.top();
        q.pop();
        num[a[tmp]]--;
        ans.pb(tmp);
        while(cnt<=Maxd&&num[cnt]==0)
        {
            for(int i=0;i<vec[cnt].size();i++)
            {
                int v=vec[cnt][i];
                q.push(v);
            }
            cnt++;
        }
    }

    for(int i=0;i<ans.size();i++)
    {
        cout<<ans[i]<<" ";
    }

    return 0;
}