Police Stations CodeForces - 796D

http://codeforces.com/group/CT2Mwjzqo4/contest/796/problem/E

有几个城市有警察局 就可以分成几棵子树

先把所有有警察局的城市扔到队列里 跑一边bfs 看没警察局的城市到哪个最近 这样就把这几个子树划分好了

 

#include <bits/stdc++.h>
using namespace std;
const int maxn=3e5+10;
const int N=0x3f3f3f3f;

struct node1
{
    int u,v;
};

struct node2
{
    int v,next;
};

queue <int> que;
node1 pre[maxn];
node2 edge[2*maxn];
int first[maxn],book[maxn],pos[maxn],belong[maxn];
int n,k,d,num;

void addedge(int u,int v)
{
    edge[num].v=v;
    edge[num].next=first[u];
    first[u]=num++;
}

void bfs()
{
    int i,u,v;
    for(i=1;i<=k;i++){
        que.push(pos[i]);
        belong[pos[i]]=pos[i];
    }
    while(!que.empty()){
        u=que.front();
        que.pop();
        for(i=first[u];i!=-1;i=edge[i].next){
            v=edge[i].v;
            if(belong[v]==0){
                que.push(v);
                belong[v]=belong[u];
            }
        }
    }
}

int main()
{
    int i,j,u,v;
    scanf("%d%d%d",&n,&k,&d);
    for(i=1,j=0;i<=k;i++){
        scanf("%d",&u);
        if(book[u]==0){
            book[u]=1;
            pos[++j]=u;
        }
    }
    k=j;
    memset(first,-1,sizeof(first));
    for(i=1;i<=n-1;i++){
        scanf("%d%d",&pre[i].u,&pre[i].v);
        addedge(pre[i].u,pre[i].v),addedge(pre[i].v,pre[i].u);
    }
    bfs();
    printf("%d\n",k-1);
    for(i=1;i<=n-1;i++) if(belong[pre[i].u]!=belong[pre[i].v]) printf("%d ",i);
    printf("\n");
    return 0;
}