# PROBLEM LINK:

*Author and Editorialist:* Dearsh Oberoi

*Tester:* Ashok Karwa

# DIFFICULTY:

Easy-Medium

# PREREQUISITES:

Sorting, BFS

# PROBLEM:

We are given an undirected connected graph with N vertices and M edges in which each node represents a city with A_i soldiers and edges bidirectional roads.

Warlords of X cities are selected and given the responsibility of choosing r more cities to recruit soldiers such that:

- Each of the r chosen cities have a path from them to the warlord that chose them.
- The cities must be nearest cities i.e. a warlord can choose a city at distance i + 1 only if all the possible cities at distance i from all warlords have already been chosen.
- The number of soldiers recruited from these cities must be maximum possible.

We have to print the maximum number of soldiers.

# QUICK EXPLANATION

First we will load all the selected X cities onto the BFS queue. Then we will perform BFS and mark the distance of cities as we come across them in an array. Declare an empty 2D vector array of size N. In it, we will store the soldiers of cities at a distance of i in the ith row. Sort all these rows in decreasing order. As we have to select the nearest cities, we will start from cities at distance 1 and go on selecting r cities.

# EXPLANATION:

Let’s assume the warlords of X selected cities to be a single unit. First, we will find the distance of all the remaining cities from this unit and store them in an array.

- To find the distances of cities, we can use BFS. We will load all the X cities onto the BFS queue and mark their distances in array as 0. The distance of the city we come across in BFS will be one greater than their parent. We will store all these distances in the index corresponding to the respective cities.

Then we will declare a 2D vector array and store the soldiers of cities at distance i in the ith row. We will sort all these rows in descending manner so that the cities with more soldiers appear early on.

As we have to select the nearest cities, we can go on by choosing the soldiers at the 1^{st} row, then the 2^{nd} row and so on of the 2D vector array we had declared.

# TIME COMPLEXITY:

*O(N + M) + NlogN*

# SOLUTIONS:

## Setter's Solution

```
#include<bits/stdc++.h>
using namespace std;
const int N = 2 * 100 * 1000 + 5;
int soldiers[N];
int depth[N];
bool visited[N];
vector< int >adj[N];
queue< int >q;
vector< int >depth_storer[N];
int max_depth = 0;
void bfs()
{
while(!q.empty())
{
int fr = q.front();
q.pop();
for(auto x : adj[fr])
{
if(!visited[x])
{
visited[x] = 1;
q.push(x);
depth[x] = depth[fr] + 1;
max_depth = max(max_depth, depth[x]);
}
}
}
}
int32_t main()
{
int n, m;
cin >> n >> m;
for(int i = 1 ; i <= n ; ++i)
cin >> soldiers[i];
for(int i = 0, a, b ; i < m ; ++i)
{
cin >> a >> b;
adj[a].push_back(b);
adj[b].push_back(a);
}
int x, r;
cin >> x >> r;
memset(visited, 0, sizeof(bool) * (n + 1));
for(int i = 0, warlord ; i < x ; ++i)
{
cin >> warlord;
visited[warlord] = 1;
q.push(warlord);
depth[warlord] = 0;
}
bfs();
for(int i = 1 ; i <= n ; ++i)
depth_storer[depth[i]].push_back(soldiers[i]);
int recruited_soldiers = 0;
for(int i = 1 ; i <= max_depth ; ++i)
{
sort(depth_storer[i].rbegin(), depth_storer[i].rend());
if(depth_storer[i].size() >= r)
{
recruited_soldiers = accumulate(depth_storer[i].begin(), depth_storer[i].begin() + r, recruited_soldiers);
break;
}
else
{
recruited_soldiers = accumulate(depth_storer[i].begin(), depth_storer[i].end(), recruited_soldiers);
r -= depth_storer[i].size();
}
}
cout << recruited_soldiers;
return 0;
}
```

## Tester's Solution

```
#include<bits/stdc++.h>
using namespace std;
int n,m;
vector<vector<int>>g;
vector<int> cnt;
vector<int> X;
vector<bool>vis;
vector<priority_queue<int>>pq;
queue<pair<int,int>>q;
void bfs(){
while(!q.empty()){
int x=q.front().second;
int l=q.front().first;
q.pop();
for(auto i:g[x]){
if(!vis[i]){
vis[i]=true;
q.push(make_pair(l+1,i));
pq[l+1].push(cnt[i]);
}
}
}
}
void solve(){
cin>>n>>m;
g.resize(n+1);
cnt.resize(n+1);
vis.resize(n+1);
for(int i=1;i<n+1;i++)
cin>>cnt[i];
for(int i=0;i<m;i++)
{
int u,v;
cin>>u>>v;
g[u].push_back(v);
g[v].push_back(u);
}
int x,r;
cin>>x>>r;
pq.resize(n+1);
for(int i=0;i<x;i++){
int u;
cin>>u;
X.push_back(u);
vis[u]=true;
}
for(int i=0;i<x;i++)
q.push(make_pair(0,X[i]));
if(x==0){
cout<<"0\n";
return;
}
bfs();
int ans=0;
int i=0;
while(r){
//cout<<i<<"-> ";
while(!pq[i].empty()){
ans+=pq[i].top();
//cout<<pq[i].top()<<" ";
pq[i].pop();
r--;
if(r==0)
break;
}
//cout<<"\n";
i++;
}
cout<<ans<<"\n";
}
int main(){
//freopen("input.txt","r",stdin); freopen("output.txt","w",stdout);
//test
{
solve();
}
return 0;
}
```