AtCoder が提供しているABC（AtCoder Beginner Contest）272 のD問題をC++とPythonで解いてみました。ABC272は、2022年10月8日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
D問題 Root M Leaper（Difficulty : 804）
問題はリンク先をご覧ください。
グラフの問題と考えてBFSを使います。辺の設定に工夫が必要です。AtCoder Problems による Difficulty は、804 でした。
解答案
左上の点からの最短距離を求める問題です。セオリー通りBFS（幅優先探索：Breadth-First search）を使います。
辺の設定には工夫が必要ですが、それはプログラムと合わせて紹介します。
C++ プログラム例（ABC272D）
点 (i1, j1) から (i2, j2) の距離を確認します。距離の2乗が $M$ であれば、辺を設定します（14、15行目）。なお、2次元の点を1次元の点に変換して処理をしています。これは以下の解説でも同じです。
残りはグラフに対するBFSの処理です。この処理は定番です。
- 距離の配列 dist と、キュー que を使います。
- キューに (0, 0) を格納して、距離 dist[0][0] = 0 とする。
- キューの要素に対して、以下を行う。
- 辺を取り出して接続している頂点までの距離が計算済でなければ、その頂点の距離を更新して、キューに格納する（30ー33行目）。
- キューが空になれば、(0, 0) から到達可能な頂点の最短距離が計算できています。
#include <bits/stdc++.h>
using namespace std;
int main()
{
int n, m;
cin >> n >> m;
vector<vector<int>> G(n * n);
for (int i1 = 0; i1 < n; ++i1) {
for (int j1 = 0; j1 < n; ++j1) {
for (int i2 = 0; i2 < n; ++i2) {
for (int j2 = 0; j2 < n; ++j2) {
if ((i1 - i2) * (i1 - i2) + (j1 - j2) * (j1 - j2) == m) {
G[i1 * n + j1].push_back(i2 * n + j2);
}
}
}
}
}
vector<int> dist(n * n, -1);
queue<int> que;
que.push(0);
dist[0] = 0;
while (!que.empty()) {
int now = que.front();
que.pop();
for (auto next : G[now]) {
if (dist[next] == -1) {
que.push(next);
dist[next] = dist[now] + 1;
}
}
}
for (int i = 0; i < n * n; ++i) {
cout << dist[i] << " \n"[i % n == n - 1];
}
return 0;
}このプログラムは入力例に対してはうまく動作しますが、提出するとTLE判定となります。辺を設定する計算量が $O(N^4)=2.56\times10^{10}$ となるためです。なんらかの工夫が必要です。
辺の設定を工夫する
実際に上記プログラムの14行目を満たすためには、$M\,-\,(i1\,-\,i2)^2$ が平方数になっている必要があります。平方数を求めるテーブルを連想配列 sq として計算しておきます（9ー14行目）。もし、$M\,-\,(i1\,-\,i2)^2 = k^2$ になっていれば、$j2 = j1 \pm k$ となり、これが $0 \leqq j2 < n$ を満たせば辺を設定します（20ー30行目）。
計算量は、$O(N^3 \log \sqrt{M}) \doteqdot 6.4 \times 10^8$ となり、余裕はないですが間に合います。
辺を設定した以降のプログラムは、前のものと同じです。
#include <bits/stdc++.h>
using namespace std;
int main()
{
int n, m;
cin >> n >> m;
map<int, int> sq;
int i = 0;
while (i * i <= m) {
sq[i * i] = i;
++i;
}
vector<vector<int>> G(n * n);
for (int i1 = 0; i1 < n; ++i1) {
for (int j1 = 0; j1 < n; ++j1) {
for (int i2 = 0; i2 < n; ++i2) {
if (sq.find(m - (i1 - i2) * (i1 - i2)) != sq.end()) {
int k = sq[m - (i1 - i2) * (i1 - i2)];
if (j1 + k < n) {
int j2 = j1 + k;
G[i1 * n + j1].push_back(i2 * n + j2);
}
if (j1 - k >= 0) {
int j2 = j1 - k;
G[i1 * n + j1].push_back(i2 * n + j2);
}
}
}
}
}
vector<int> dist(n * n, -1);
queue<int> que;
que.push(0);
dist[0] = 0;
while (!que.empty()) {
int now = que.front();
que.pop();
for (auto next : G[now]) {
if (dist[next] == -1) {
que.push(next);
dist[next] = dist[now] + 1;
}
}
}
for (int i = 0; i < n * n; ++i) {
cout << dist[i] << " \n"[i % n == n - 1];
}
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC272D）
Python 版も基本的な考え方は同じです。キューは deque を使います。
以下のプログラムは、「Python (CPython 3.11.4)」では、TLE判定となり、「Python (PyPy 3.10-v7.3.12)」で、AC判定となりました。
"""AtCoder Beginner Contest 272 D"""
from collections import deque
n, m = map(int, input().split())
sq = {}
i = 0
while i * i <= m:
sq[i * i] = i
i += 1
G = [[] for _ in range(n * n)]
for i1 in range(n):
for j1 in range(n):
for i2 in range(n):
if m - (i1 - i2) ** 2 in sq:
k = sq[m - (i1 - i2) ** 2]
if j1 + k < n:
j2 = j1 + k
G[i1 * n + j1].append(i2 * n + j2)
if j1 - k >= 0:
j2 = j1 - k
G[i1 * n + j1].append(i2 * n + j2)
dist = [-1] * (n * n)
que = deque()
que.append(0)
dist[0] = 0
while que:
now = que.popleft()
for next_v in G[now]:
if dist[next_v] == -1:
que.append(next_v)
dist[next_v] = dist[now] + 1
for i in range(n * n):
print(dist[i], end=" " if i % n != n - 1 else "\n")最後に
ABC272は、ABCを再開した少し後に参加しました。この問題は、まったく解けませんでした。
実力を上げるために前に解けなかった問題を解いています。この問題は、いまでは慣れたBFSを使い、辺の張り方で工夫して解くことができました。
引き続き ABC の問題を紹介していきます。
