[CF796D] Police Stations

Description

给定一棵树，树上有一些点是关键点。一棵树是合法的当且仅当所有点离最近的关键点的距离不超过 (d)，求最多能删除多少条边。

Solution

考虑从所有的关键点开始 BFS。假设当前位于点 (p)，下一个要访问的点是 (q)，若 (q) 被访问过而 ((p,q)) 这条边未被访问过，则这条边可以被删除。

#include <bits/stdc++.h>
using namespace std;

#define int long long 
const int N = 1000005;

int n;
int k;                          // 关键点个数
int d;                          // 距离限制
int vis[N];                     // 点是否被访问过
unordered_set <int> eVis;       // 边是否被访问过
unordered_set <int> eSrc;       // E
unordered_set <int> eAns;       // 可以删除的边的集合
unordered_map <int,int> eInd;   // 边的序号
vector <int> keyVertex;         // 关键点

int edgeId(int u,int v)         // 获取边的编码
{
    if(u>v) swap(u,v);
    return u*N+v;
}

void setEdgeVis(int u,int v) 
{
    eVis.insert(edgeId(u,v));
}

int getEdgeVis(int u,int v)
{
    if(eVis.find(edgeId(u,v))==eVis.end()) return false;
    return true;
}

vector <int> g[N];

struct BFSData      // BFS 队列中的信息
{
    int p;          // 当前点下标
    int dist;       // 到出发点的距离
};

void bfs()
{
    queue <BFSData> que;
    for(auto p:keyVertex)
    {
        que.push({p,0});
        vis[p]=1;
    }
    while(!que.empty())
    {
        int p=que.front().p;
        int dist=que.front().dist;
        que.pop();
        if(dist==d) continue;
        for(auto q:g[p]) 
        {
            if(vis[q] && !getEdgeVis(p,q))
            {
                eAns.insert(edgeId(p,q));
            }
            if(!vis[q])
            {
                que.push({q,dist+1});
                setEdgeVis(p,q);
                vis[q]=1;
            }
        }
    }
}

signed main()
{
    cin>>n>>k>>d;
    for(int i=1;i<=k;i++)
    {
        int tmp;
        cin>>tmp;
        keyVertex.push_back(tmp);
    }
    for(int i=1;i<n;i++)
    {
        int u,v;
        cin>>u>>v;
        g[u].push_back(v);
        g[v].push_back(u);
        eInd[edgeId(u,v)]=i;
        eSrc.insert(edgeId(u,v));
    }

    bfs();

    for(auto i:eSrc)
    {
        if(eVis.find(i)==eVis.end()) 
        {
            eAns.insert(i);
        }
    }

    cout<<eAns.size()<<endl;
    for(auto i:eAns)
    {
        cout<<eInd[i]<<" ";
    }
    return 0;
}