ABC254-E 解説

E - Small d and k

各頂点の次数は $3$ 以下、そして各クエリで $k_i$ は $3$ 以下という制約があります。すると、各クエリで見る必要がある頂点は $40$ 頂点に抑えることができるので、BFSをすることで解けます。ただし、訪問の判定に、毎回全ての頂点に対して訪れたかを持つのは無理なので、PythonのdefaultdictやC++のmapで距離を管理します。

from collections import defaultdict, deque
N, M = map(int, input().split())
g = [[] for _ in range(N)]
for _ in range(M):
    u, v = map(int, input().split())
    u -= 1
    v -= 1
    g[u].append(v)
    g[v].append(u)
for _ in range(int(input())):
    x, k = map(int, input().split())
    x -= 1
    d = defaultdict(lambda: -1)
    d[x] = 0
    q = deque([x])
    while q:
        u = q.popleft()
        if d[u] < k:
            for v in g[u]:
                if d[v] == -1:
                    d[v] = d[u] + 1
                    q.append(v)
    ans = 0
    for i in d.keys():
        ans += i + 1
    print(ans)

以下のC++の実装ではmapのデフォルトの値を-1にする方法がわからないので、距離+1を管理しています。

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

int main() {
    int N, M;
    cin >> N >> M;
    vector<vector<int>> g(N);
    for (int i = 0; i < M; i++) {
        int u, v;
        cin >> u >> v;
        u--; v--;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    int Q;
    cin >> Q;
    for (int i = 0; i < Q; i++) {
        int x, k;
        cin >> x >> k;
        x--;
        map<int, int> d;
        d[x] = 1;
        deque<int> q;
        q.push_back(x);
        while (!q.empty()) {
            int u = q.front();
            q.pop_front();
            if (d[u] <= k) {
                for (int v: g[u]) {
                    if (d[v] == 0) {
                        d[v] = d[u] + 1;
                        q.push_back(v);
                    }
                }
            }
        }
        int ans = 0;
        for (auto [i, j]: d) {
            ans += i + 1;
        }
        cout << ans << '\n';
    }
}