ABC254E - Small d and k

考え方

制約から

• 距離0：1コ
• 距離1：1×3コ
• 距離2：1×3×3コ
• 距離3：1×3×3×3コ

で，各クエリに対して40コ調べれば良いから，言われたとおり計算するだけでよい．

回答例

訪れた頂点の管理をするために全頂点の配列を用意するとTLEする（各頂点への距離distを配列で用意するとか，全頂点にたいし訪れたかどうかをTrue / False で管理する配列を用意するとか）．

配列ではなくsetを使うと間に合う．

from collections import deque
import sys
input = sys.stdin.readline

N, M = map(int, input().split())
G = [[] for _ in range(N)]

for _ in range(M):
    a, b = map(lambda x: int(x) - 1, input().split())
    G[a].append(b)
    G[b].append(a)

Q = int(input())
for _ in range(Q):
    x, k = map(int, input().split())
    seen = set([x - 1])
    q = deque([(x - 1, 0)])
    ans = x
    while q:
        cur, d = q.popleft()
        if d + 1 > k: break
        for chi in G[cur]:
            if chi in seen: continue
            seen.add(chi)
            q.append((chi, d + 1))
            ans += chi + 1
    print(ans)

さらに

import sys
input = sys.stdin.readline

を加えると，もう少し早くなる．

queに入れるタイミングではなく，popするタイミングでansに加算しても良い．

from collections import deque
import sys
input = sys.stdin.readline

N, M = map(int, input().split())
G = [[] for _ in range(N)]

for _ in range(M):
    a, b = map(lambda x: int(x) - 1, input().split())
    G[a].append(b)
    G[b].append(a)

Q = int(input())
for _ in range(Q):
    x, k = map(int, input().split())
    seen = set([x - 1])
    q = deque([(x - 1, 0)])
    ans = 0
    while q:
        cur, d = q.popleft()
        if d <= k:
            ans += cur + 1
        if d == k:
            continue
        for chi in G[cur]:
            if chi in seen: continue
            seen.add(chi)
            q.append((chi, d + 1))
    print(ans)

【参考】Submission #32218930 - AtCoder Beginner Contest 254