Official

E - Small d and k Editorial by m_99


\(\mathrm O (MQ)\) 解法

クエリごとに幅優先探索(BFS)を行って頂点 \(x_i\) から各頂点までの距離を求め、距離が \(k_i\) 以下の頂点の番号の総和を求めれば答え自体は正しく求められます。しかし、BFS \(1\) 回につき \(\mathrm O(M)\) であるため、全体で \(\mathrm O (MQ)\) となり遅すぎます。

条件を満たす頂点についての考察

以下のように計算をすると、頂点 \(x_i\) からの距離が \(k_i(\leq 3)\) 以下である頂点の個数は少ないことが分かります。

  • 距離が \(0\) の頂点は頂点 \(x_i\) のみ
  • 距離が \(1\) の頂点は頂点 \(x_i\) に隣接する頂点。次数が \(3\) 以下なので高々 \(3\) 個。
  • 距離が \(2\) の頂点は距離が \(1\) の頂点に隣接する頂点(の一部)。距離が \(1\) の頂点が \(3\) 個、次数が \(3\) 以下なので高々 \(3 \times 3 = 9\) 個。
  • 距離が \(3\) の頂点は距離が \(2\) の頂点に隣接する頂点(の一部)。距離が \(2\) の頂点が \(9\) 個、次数が \(3\) 以下なので高々 \(9 \times 3 = 27\) 個。
  • 以上を合計すると \(40\)

なお、次数が \(3\) 以下ということから距離が \(k\) の頂点の個数を距離が \(k-1\) の頂点の個数の \(3\) 倍以下としていますが、距離が \(2\) 以上の頂点に関しては\(3\) 本の辺のうち少なくとも \(1\) 本は距離がより短い頂点から来るのに使われるため \(2\) 倍以下と評価することもできます。
いずれにせよ、条件を満たす頂点のみを列挙して番号の総和を求める分には \(Q\) 回繰り返してもTLEの心配がありません。

BFS・DFSで条件を満たす頂点を列挙

以上より、BFSを距離が \(k_i\) 以下の部分のみ行う、つまり距離が \(k_i\) の頂点まで調べたらそれ以上調べないことにすれば本問題を解くことが出来ます。詳しくは実装例を参考にしてください。(距離配列の一部のみを使うために訪れた頂点をvsとして記録する等していますが、setやmapを使う方が簡便かもしれません)
また、上述の評価を考えると深さ優先探索(DFS)でも十分高速です。

実装例

#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: