Official

E - Small d and k Editorial by en_translator


\(\mathrm O (MQ)\) Solution

We may obtain the correct answer itself for each query by performing BFS (Breadth-First Search) to find the distance from Vertex \(x_i\) to each vertex and computing the sum of indices of vertex whose distance is at most \(k_i\). However, since each BFS costs \(O(M)\) time, it costs a total of \(\mathrm O (MQ)\), which is too slow.

Observation on vertices satisfying the condition

We can see from the calculation below that there are not so many vertices whose distance from Vertex \(x_i\) is at most \(k_i(\leq 3)\).

  • Only Vertex \(xi\) has the distance \(0\).
  • Only the vertices adjacent to Vertex \(x_i\) has the distance \(1\). Since its degree is at most \(3\), there are at most \(3\) such vertices.
  • Only (some of) the vertices adjacent to vertices with the distance \(1\) have the distance \(2\). Since there are at most \(3\) vertices, each of degree at most \(3\), there are at most \(3 \times 3 = 9\) such vertices.
  • Only (some of) the vertices adjacent to vertices with the distance \(2\) have the distance \(3\). Since there are at most \(9\) vertices, each of degree at most \(3\), there are at most \(9 \times 3 = 27\) such vertices.
  • The sum of the vertices above is \(40\).

In the discussion above, we evaluated the upper bound of the number of vertices with the distance \(k\) by \(3\) times the number of vertices with the distance \(k-1\), but we can replace “\(3\) times” with “\(2\) times” because among the three edges adjacent to a vertex of distance at least \(2\), at least one of them comes from a vertex with smaller distance.
Anyway, we may safely repeat enumerating the indices of the vertices satisfying the condition and find their times \(Q\) times without TLE.

Enumerating the vertices satisfying the condition with BFS or DFS

Thus, this problem can be solved by performing BFS only on vertices with distances at most \(k_i\), i.e. stopping BFS once you reached a vertex with distance \(k_i\). For more details, please refer to the Sample code. (We store the visited vertices to vs in order to use only some elements of the distance array, but it may be easier to use a set or a map)
Also, by the evaluation above, Depth-First Search (DFS) is also sufficiently fast.

Sample code

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

int main() {
    
	int N,M;
	cin>>N>>M;
	vector<vector<int>> E(N);
	for(int i=0;i<M;i++){
		int a,b;
		cin>>a>>b;
		a--,b--;
		E[a].push_back(b);
		E[b].push_back(a);
	}
	
	vector<int> dis(N,-1);
	
	int Q;
	cin>>Q;
	
	for(int i=0;i<Q;i++){
		int x,k;
		cin>>x>>k;
		x--;
		queue<int> q;
		q.push(x);
		dis[x] = 0;
		vector<int> vs;
		while(q.size()>0){
			int u = q.front();
			q.pop();
			vs.push_back(u);
			if(dis[u]==k){
				continue;
			}
			for(int j=0;j<E[u].size();j++){
				int v = E[u][j];
				if(dis[v]!=-1)continue;
				dis[v] = dis[u] + 1;
				q.push(v);
			}
		}
		int ans = 0;
		for(int j=0;j<vs.size();j++){
			ans += vs[j]+1;
			dis[vs[j]] = -1;
		}
		cout<<ans<<endl;
	}
	
    return 0;
}

posted:
last update: