PyPyとPythonの速度比較を兼ねて。解説にDFS や BFS, Union-Findが例示されているのでそれぞれのコード例です。
参考にダイクストラ法やBellman-Ford法で解こうとするとどうなるの？もやってみました。
| アプローチ | 実行時間(pypy3) | submit |
|---|---|---|
| DFS(再帰-Pypy) | 136ms | Submit |
| DFS(deque-Pypy) | 133ms | Submit |
| BFS(deque) | 139ms | Submit |
| Union-Find | 180ms | Submit |
| ダイクストラ法 | 536ms | Submit |
| Bellman-Ford法 | TLE | Submit |
| アプローチ | 実行時間(Python) | submit |
|---|---|---|
| DFS(再帰-Python) | 1351ms | Submit |
| DFS(deque-Python) | 1589ms | Submit |
| BFS(deque-Python) | 1627 ms | Submit |
| Union-Find(Python) | TLE | Submit |
| ダイクストラ法(Python) | TLE | Submit |
| Bellman-Ford法(Python) | TLE | Submit |
DFS(深さ優先探索 - 再帰する)
1のノードから探索したことのないのノードを探索していきます。
import sys
sys.setrecursionlimit(100000)
n, d = map(int, input().split())
dat = []
for _ in range(n): dat.append(tuple(map(int, input().split())))
visited = [False] * n
def func(cur):
x1, y1 = dat[cur]
for i in range(n): # 全探索
if visited[i]: continue # 既に感染
x2, y2 = dat[i]
if (x1-x2)**2 + (y1-y2)**2 <= (d**2):
visited[i] = True
func(i)
visited[0] = True
func(0)
for x in visited: print("Yes" if x else "No")
DFS(深さ優先探索 - 非再帰)
1のノードから探索したことのないのノードをスタックに入れていけばよいです。dequeを使うならpopleftしながらappendleftしていきます。
# DFS
from collections import deque
n, d = map(int, input().split())
dat = []
for _ in range(n): dat.append(tuple(map(int, input().split())))
visited = [False] * n
infection = deque([0])
visited[0] = True
while len(infection) > 0:
cur = infection.popleft()
x1, y1 = dat[cur]
for i in range(n): # 全探索
if visited[i]: continue # 既に感染
x2, y2 = dat[i]
if (x1-x2)**2 + (y1-y2)**2 <= (d**2):
infection.appendleft(i) # leftに入るのでDFSになる。どんどん掘る。
visited[i] = True
for x in visited: print("Yes" if x else "No")
BFS(幅優先探索)
1のノードから探索したことのないのノードをキューに入れていけばよいです。dequeを使うならpopleftしながらappendしていきます。
from collections import deque
n, d = map(int, input().split())
dat = []
for _ in range(n): dat.append(tuple(map(int, input().split())))
visited = [False] * n
infection = deque([0])
visited[0] = True
while len(infection) > 0:
cur = infection.popleft()
x1, y1 = dat[cur]
for i in range(n): # 全探索
if visited[i]: continue # 既に感染
x2, y2 = dat[i]
if (x1-x2)**2 + (y1-y2)**2 <= (d**2):
infection.append(i) # rightに入るのでBFSになる
visited[i] = True
for x in visited: print("Yes" if x else "No")
Union-Find
not522さんのライブラリを使っています。全点のペアを全探索してmergeしていき、1のノードと同じ連結成分(same)かを出力すればいいです。
# DSUライブラリ部分は省略しています
from collections import deque
n, d = map(int, input().split())
dat = []
for _ in range(n): dat.append(tuple(map(int, input().split())))
dsu = DSU(n)
for i in range(n):
x1, y1 = dat[i]
for j in range(i+1, n):
x2, y2 = dat[j]
if (x1-x2)**2 + (y1-y2)**2 <= (d**2):
dsu.merge(i, j)
for i in range(n): print("Yes" if dsu.same(0, i) else "No")
おまけ：ダイクストラ法
この問題はノード1から到達可能なノードを調べる問題なので、単始点から全点への到達可能な最短距離を調べるダイクストラ法を使うこともできそうです。「ライブラリを張るだけ！距離も分かる！」...ではあるのですが、最短距離を調べる必要がないのでtoo muchです。
辺を張る動作はDSUでmergeするのと同じ操作でほぼメインのコードの形は変わりません。距離も分かりますがBFSでも最短距離は分かります。
# dijkstra部分は省略
n, d = map(int, input().split())
dat = []
for _ in range(n): dat.append(tuple(map(int, input().split())))
dj = dijkstra(n)
for i in range(n):
x1, y1 = dat[i]
for j in range(i+1, n):
x2, y2 = dat[j]
if (x1-x2)**2 + (y1-y2)**2 <= (d**2):
dj.makeEdge(i, j, 1)
cost = dj.solve(0)
for i in range(n): print("Yes" if dj.cost[i] != dj.INF else "No")
おまけ2：Bellman-Ford(TLE)
ダイクストラ法が使えるのならBellman-Ford法でもいいのでは？と思っても間に合いません。今回は完全グラフになるため、辺の数は$N^2$のため、Bellman-Ford法の計算量$O(|V| |E|)$より$O(N^3)$となり間に合いませんでした。