AtCoder が提供しているABC(AtCoder Beginner Contest)348 のD問題をC++とPythonで解いてみました。ABC348は、2024年4月6日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
D問題 Medicines on Grid(Difficulty : 1108)
問題はリンク先をご覧ください。
スタートとゴール、エネルギーがある場所をグラフの頂点と考えて、連結しているか調べます。AtCoder Problems による Difficulty は 1108 でした。
解答案
スタート地点、ゴール地点、エネルギー(薬)がある地点をグラフの頂点と考えます。グラフの辺が設定できれば、スタート地点とゴール地点が連結しているか調べるだけです。
エネルギーがある地点から、グリッドの最短距離をBFSで求めることができます。最短距離がエネルギー以下であれば、エネルギーがある地点から到達できます。この範囲に、他のグラフの頂点があれば、その頂点との間に有向辺を設定します(向きがあることに注意)。
C++ プログラム例(ABC348D)
プログラムの補足です。
- スタート地点とゴール地点にエネルギーがあるか調べます(31ー43行目)。
- スタート地点にエネルギーがない場合は、”No” を出力して終了します(44ー47行目)。
- ゴール地点にエネルギーがない場合は、グラフの頂点の候補にゴール地点を加えます。持っエネルギーは0とします(48ー60行目)。
- グリッドの最短距離を求める定番コードです(71ー90行目)。
- 最短距離が求まれば、頂点間の有向辺を設定します。わたしの定番コードは、距離の初期値を-1としていたため、条件(96行目)に注意が必要です(92ー99行目)。
- グラフの頂点と辺が設定できれば、連結性を調べます。慣れているDFSを使いました(9ー17行目、103行目)。
長くなりましたが、以下がC++プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef pair<int, int> P;
vector<vector<int>> G;
vector<bool> seen;
void dfs(int v)
{
seen[v] = true;
for (auto next_v : G[v]) {
if (!seen[next_v]) {
dfs(next_v);
}
}
}
int main()
{
int h, w;
cin >> h >> w;
vector<string> s(h);
for (int i = 0; i < h; ++i) {
cin >> s[i];
}
int n;
cin >> n;
vector<int> r(n), c(n), e(n);
int start_v = -1;
int end_v = -1;
for (int i = 0; i < n; ++i) {
cin >> r[i] >> c[i] >> e[i];
--r[i];
--c[i];
if (s[r[i]][c[i]] == 'S') {
start_v = i;
}
if (s[r[i]][c[i]] == 'T') {
end_v = i;
}
}
if (start_v == -1) {
cout << "No" << endl;
return 0;
}
if (end_v == -1) {
for (int i = 0; i < h; ++i) {
for (int j = 0; j < w; ++j) {
if (s[i][j] == 'T') {
r.push_back(i);
c.push_back(j);
e.push_back(0);
n = r.size();
end_v = n - 1;
}
}
}
}
G.resize(n);
int di[4] = {1, 0, -1, 0};
int dj[4] = {0, 1, 0, -1};
for (int i = 0; i < n; ++i) {
queue<P> que;
vector<vector<int>> dist(h, vector<int>(w, -1));
que.push(P(r[i], c[i]));
dist[r[i]][c[i]] = 0;
while (!que.empty()) {
P now = que.front();
que.pop();
int now_i = now.first;
int now_j = now.second;
for (int d = 0; d < 4; ++d) {
int next_i = now_i + di[d];
int next_j = now_j + dj[d];
if ((0 <= next_i)&&(next_i < h)
&&(0 <= next_j)&&(next_j < w)
&&(dist[next_i][next_j] == -1)
&&(s[next_i][next_j] != '#')) {
que.push(P(next_i, next_j));
dist[next_i][next_j] = dist[now_i][now_j] + 1;
}
}
}
for (int j = 0; j < n; ++j) {
if (i == j) {
continue;
}
if ((dist[r[j]][c[j]] != -1)&&(dist[r[j]][c[j]] <= e[i])) {
G[i].push_back(j);
}
}
}
seen.assign(n, false);
dfs(start_v);
if (seen[end_v]) {
cout << "Yes" << endl;
} else {
cout << "No" << endl;
}
return 0;
}
グラフとして設定できたら、スタート地点とゴール地点の連結性を調べるためにBFSを使うこともできます。こちらも定番コードです(89ー102行目)。
こちらが公式解説の想定解法(解法1)のようです。
#include <bits/stdc++.h>
using namespace std;
typedef pair<int, int> P;
int main()
{
int h, w;
cin >> h >> w;
vector<string> s(h);
for (int i = 0; i < h; ++i) {
cin >> s[i];
}
int n;
cin >> n;
vector<int> r(n), c(n), e(n);
int start_v = -1;
int end_v = -1;
for (int i = 0; i < n; ++i) {
cin >> r[i] >> c[i] >> e[i];
--r[i];
--c[i];
if (s[r[i]][c[i]] == 'S') {
start_v = i;
}
if (s[r[i]][c[i]] == 'T') {
end_v = i;
}
}
if (start_v == -1) {
cout << "No" << endl;
return 0;
}
if (end_v == -1) {
for (int i = 0; i < h; ++i) {
for (int j = 0; j < w; ++j) {
if (s[i][j] == 'T') {
r.push_back(i);
c.push_back(j);
e.push_back(0);
n = r.size();
end_v = n - 1;
}
}
}
}
vector<vector<int>> G(n);
int di[4] = {1, 0, -1, 0};
int dj[4] = {0, 1, 0, -1};
for (int i = 0; i < n; ++i) {
queue<P> que;
vector<vector<int>> dist(h, vector<int>(w, -1));
que.push(P(r[i], c[i]));
dist[r[i]][c[i]] = 0;
while (!que.empty()) {
P now = que.front();
que.pop();
int now_i = now.first;
int now_j = now.second;
for (int d = 0; d < 4; ++d) {
int next_i = now_i + di[d];
int next_j = now_j + dj[d];
if ((0 <= next_i)&&(next_i < h)
&&(0 <= next_j)&&(next_j < w)
&&(dist[next_i][next_j] == -1)
&&(s[next_i][next_j] != '#')) {
que.push(P(next_i, next_j));
dist[next_i][next_j] = dist[now_i][now_j] + 1;
}
}
}
for (int j = 0; j < n; ++j) {
if (i == j) {
continue;
}
if ((dist[r[j]][c[j]] != -1)&&(dist[r[j]][c[j]] <= e[i])) {
G[i].push_back(j);
}
}
}
vector<bool> seen(n, false);
queue<int> que;
seen[start_v] = true;
que.push(start_v);
while (!que.empty()) {
int now = que.front();
que.pop();
for (auto next_v : G[now]) {
if (!seen[next_v]) {
seen[next_v] = true;
que.push(next_v);
}
}
}
if (seen[end_v]) {
cout << "Yes" << endl;
} else {
cout << "No" << endl;
}
return 0;
}
AC(Accepted=正しいプログラム)と判定されました。
Python プログラム例(ABC348D)
Python は、再帰関数の呼び出し回数に制限があり速くもないようです。このため、BFSで連結性を調べました。キューは collections.deque を使いました。
以下のプログラムは、「Python (CPython 3.11.4)」では、TLE判定となり、「Python (PyPy 3.10-v7.3.12)」で、AC判定となりました。
"""AtCoder Beginner Contest 348 D"""
import sys
from collections import deque
h, w = map(int, input().split())
s = [input() for i in range(h)]
n = int(input())
r = [0] * n
c = [0] * n
e = [0] * n
start_v = -1
end_v = -1
for i in range(n):
r[i], c[i], e[i] = map(int, input().split())
r[i] -= 1
c[i] -= 1
if s[r[i]][c[i]] == 'S':
start_v = i
if s[r[i]][c[i]] == 'T':
end_v = i
if start_v == -1:
print("No")
sys.exit()
if end_v == -1:
for i in range(h):
for j in range(w):
if s[i][j] == 'T':
r.append(i)
c.append(j)
e.append(0)
n = len(r)
end_v = n - 1
G = [[] for _ in range(n)]
di = [1, 0, -1, 0]
dj = [0, 1, 0, -1]
for i in range(n):
que = deque()
dist = [[-1 for _ in range(w)] for _ in range(h)]
que.append((r[i], c[i]))
dist[r[i]][c[i]] = 0
while que:
now_i, now_j = que.popleft()
for d in range(4):
next_i = now_i + di[d]
next_j = now_j + dj[d]
if 0 <= next_i < h and 0 <= next_j < w and \
dist[next_i][next_j] == -1 and s[next_i][next_j] != '#':
que.append((next_i, next_j))
dist[next_i][next_j] = dist[now_i][now_j] + 1
for j in range(n):
if i == j:
continue
if dist[r[j]][c[j]] != -1 and dist[r[j]][c[j]] <= e[i]:
G[i].append(j)
seen = [False] * n
que = deque()
seen[start_v] = True
que.append(start_v)
while que:
now = que.popleft()
for next_v in G[now]:
if not seen[next_v]:
seen[next_v] = True
que.append(next_v)
print("Yes" if seen[end_v] else "No")
こちらも「AC」と判定されました。
最後に
今回は、バーチャルコンテストの参加でした。この問題は自分で知っている範囲の知識で解ける問題ですが、実装にとても時間がかかりました(コンテスト時間ではWAが取れなかったです)。これからも問題を解くことを楽しんでいきます。
引き続き ABC の問題を紹介していきます。