AtCoder が提供しているABC（AtCoder Beginner Contest）324 のD問題をC++とPythonで解いてみました。ABC324は、2023年10月14日21:00に実施されました。
AtCoder の紹介はこちらに、プログラミングの方針はこちらに記事があります。
D問題 Square Permutation（Difficulty : 895）
問題はリンク先をご覧ください。
ある桁数を持つ平方数について調べます。気づきにくいコーナーケースがあります。AtCoder Problems による Difficulty は、895 でした。
解答案
$N$ 桁の平方数を求めて、条件を満たすか確認します。注意点としては、文字列 $S$ に ’0’ が含まれている場合です。仮に $N=4$ の場合、平方数が 1000 以上、10000未満の場合を調べればよいですが、文字列に ‘0’ を一つ含む場合、もう1桁下の 100 以上の場合を調べる必要があります。
入力例1 “4320” は、$N=4$ 桁ですので、1000以上、10000未満の場合を調べればよいと思えますが、千の桁が ‘0’ である場合も調べる必要があります。このため 100 以上の場合も調べる必要があります。実際に 0324=324=182 は平方数になり該当します。
この問題には、ひっかかりやすいコーナーケースがあります。問題には明記されていませんが、02 = 0 なので 0 は平方数になります。実際にこの場合を確かめるテストケースがあります（03_small_36.txt）。
C++ プログラム例（ABC324D）
プログラムのコメントは以下となります。
- 文字列 T が含む数字を配列 sd に格納します（13ー16行目）。
- 文字列 T が ‘0’ しか含まれないコーナーケースを処理します（18ー21行目）。
- $N$ と文字列に含まれる ‘0’ の個数から、平方数の範囲 [start, end) を求めておきます（23ー30行目）。
- 小さな順に平方数を調べていきます。
- 平方数が [start, end) に含まれている場合、含まれている数字を調べて sd と比較します。桁数が足りない場合は、’0′ の数を加えます（50行目）。
- 含まれている数字が一致した場合は、変数 result をインクリメントします。
以下がC＋＋プログラムとなります。
#include <bits/stdc++.h>
using namespace std;
typedef unsigned long long int ull;
int main()
{
int n;
cin >> n;
string s;
cin >> s;
vector<int> sd(10, 0);
for (int i = 0; i < n; ++i) {
++sd[s[i] - '0'];
}
if (sd[0] == n) {
cout << 1 << endl;
return 0;
}
ull start = 1;
ull end = 1;
for (int i = 0; i < n - 1 - sd[0]; ++i) {
start *= 10;
}
for (int i = 0; i < n; ++i) {
end *= 10;
}
ull result = 0;
ull i = 0;
while (1) {
++i;
ull num = i * i;
if (num >= end) {
break;
}
if (num < start) {
continue;
}
vector<int> d(10, 0);
int keta = 0;
while (num > 0) {
++keta;
++d[num % 10];
num /= 10;
}
d[0] += n - keta;
if (sd == d) {
++result;
}
}
cout << result << endl;
return 0;
}AC（Accepted＝正しいプログラム）と判定されました。
Python プログラム例（ABC324D）
Python 版も基本的な考え方は同じです。
以下のプログラムは、「Python (CPython 3.11.4)」では、TLE判定となり、「Python (PyPy 3.10-v7.3.12)」で、AC判定となりました。
"""AtCoder Beginner Contest 324 D"""
import sys
n = int(input())
s = input()
sd = [0] * 10
for _, ch in enumerate(s):
sd[ord(ch) - ord('0')] += 1
if sd[0] == n:
print(1)
sys.exit(0)
start = 1
for i in range(n - 1 - sd[0]):
start *= 10
end = 10 ** n
result = 0
i = 0
while True:
i += 1
num = i * i
if num >= end:
break
if num < start:
continue
d = [0] * 10
keta = 0
while num > 0:
keta += 1
d[num % 10] += 1
num //= 10
d[0] += n - keta
if sd == d:
result += 1
print(result)最後に
コンテストでは、解きやすいと感じたE問題を先に解きました。結果的にE問題はACを取れましたが、この問題は、時間切れでした。全体的に解くスピードを上げていく必要があるようです。
引き続き ABC の問題を紹介していきます。
