AtCoder

ABC272 D問題(Root M Leaper)を解く

AtCoder_ABC272_D

AtCoder が提供しているABC(AtCoder Beginner Contest)272 のD問題をC++とPythonで解いてみました。ABC272は、2022年10月8日21:00に実施されました。

AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。

D問題 Root M Leaper(Difficulty : 804)

問題はリンク先をご覧ください。

ABC272 D問題 Root M Leaper

グラフの問題と考えて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 の問題を紹介していきます。

COMMENT

メールアドレスが公開されることはありません。 が付いている欄は必須項目です

CAPTCHA