精進（2024/4/19） - あるぱかのブログ

最近のABCからD問題を解く

HUAWEI Programming Contest 2024（AtCoder Beginner Contest 342） D - Square Pair
トヨタ自動車プログラミングコンテスト2024#4（AtCoder Beginner Contest 348） D - Medicines on Grid

HUAWEI Programming Contest 2024（AtCoder Beginner Contest 342） D - Square Pair

概要：配列 $A$ の中から2つ選んで掛け算したときに積が平方数になる組み合わせ数を数えてね

$N \leq 2 \times 10^{5}$ から2つ選ぶ組み合わせ数は $_N C _2 = \dfrac{N \times (N-1)}{2}$ で最大 $2 \times 10^{12}$ になるので愚直にやるとTLEになります。

本番では全く見当もつかなかったので解説で方針だけ確認しました。

配列 $A$ の各要素についてこれを割り切る最大の平方数 $d_{A_i}$ を求める（ $A_i$ =0 のときは0とする）
$\dfrac{A_i}{d_{A_i}}$ が同じになる組み合わせ数を数え上げる

$\dfrac{A_i}{d_{A_i}} \dfrac{A_j}{d_{A_j}}$ が平方数であることの証明は記載ありませんでしたが、おそらくこういうことだと思います。

$\dfrac{A_i}{d_{A_i}}$ の約数に平方数が含まれない
約数に平方数を持たない数同士を掛けて平方数にするには両者が同じ数字でなければならない

この方針により問題を配列 $D$ （ $D_i = \dfrac{A_i}{d_{A_i}}$ ）の度数分布と組み合わせ数の計算に変換でき、計算量を $O(N \sqrt(A_i))$ にできました。

$A$ に0があるときが厄介で、0に対しては何を掛けても0になり題意を満たしてしまうので特別な処理が必要になります。

（0の処理で3回くらいWAしました・・・）

#ABC342 D - Square Pair
#解説
from collections import Counter
N = int(input())
A = list(map(int, input().split()))
D = [0]*N
z = 0
for i in range(N):
    #Ai=0なら0を返す
    if A[i] == 0:
        D[i] = 0
        z += 1
    #それ以外は割り切れる最大の平方数で割った商を返す
    else:
        #探索範囲は√AiまででOK
        for j in range(1, int(A[i]**0.5)+1):
            if A[i] % j**2 == 0:
                d = j**2
        D[i] = A[i] // d
ans = 0
num = Counter(sorted(D))
for key in num:
    if key == 0:
        for i in range(1, num[key]+1):
            ans += N-i
    else:
        ans += (num[key] * (num[key] - 1)) // 2
print(ans)

提出リンク：提出 #52513574 - HUAWEI Programming Contest 2024（AtCoder Beginner Contest 342）

トヨタ自動車 プログラミングコンテスト2024#4（AtCoder Beginner Contest 348） D - Medicines on Grid

リンク：D - Medicines on Grid (atcoder.jp)

概要：エネルギー管理しながら二次元盤面のゴールにたどり着けるか判定してね

BFSのいい練習になりました。

これを機に二次元盤面とグラフのBFSのテンプレートを作ったので次からは安心！

公式解説の解法1の通り以下の方針で解きます。

スタートおよび薬のあるマスを始点としてBFS、エネルギーの続くうちに他の薬orゴールに到達可能か判定して、可能なら有効グラフに辺を追加する
上記のグラフでBFSしてスタートからゴールが続いているか判定して答えを出力する

私がやった間違いとして

二次元盤面探索で始点と終点を指定して到達可能か判定するように関数を作ってしまったので計算量が $O(HWN^{2})$ になってしまいTLE
（提出リンク：提出 #52248291 - トヨタ自動車プログラミングコンテスト2024#4（AtCoder Beginner Contest 348）
⇒始点だけ指定して各マスへの最短距離を書いた二次元配列 $dist$ を返して薬マスに到達可能か（ $dist_{薬i} \geq 0$ か）見ることで $O(HWN)$ に改善

初めて100行近いコード書きました。疲れた・・・

 from collections import deque
#二次元盤面BFSで各マスへの最短距離を書いた二次元配列distを返す
def BFS_2Dgrid(grid,start):
    H = len(grid)
    W = len(grid[0])
    DX = [1,-1,0,0]
    DY = [0,0,1,-1]
    sx, sy = start
    gx, gy = goal
    q = deque()
    q.append(start)
    dist = [[-1 for i in range(W)] for j in range(H)]
    dist[sx][sy] = 0
    while q:
        x, y = q.popleft()
        for dx, dy in zip(DX,DY):
            xt, yt = x+dx, y+dy
            if xt < 0 or H-1 < xt or yt < 0 or W-1 < yt:
                continue
            if grid[xt][yt] == '#':
                continue
            if dist[xt][yt] > -1:
                continue
            else:
                dist[xt][yt] = dist[x][y] + 1
                q.append((xt,yt))
    return dist
#1次元グラフ上のBFSでゴールへの最短距離を返す
def BFS_1Dgraph(graph,start,goal):
    N = len(graph)
    q = deque()
    q.append(start)
    dist = [-1 for i in range(N)]
    dist[start] = 0
    while q:
        pop = q.popleft()
        if pop == goal:
            break
        else:
            for i in G[pop]:
                if dist[i] > -1:
                    continue
                else:
                    dist[i] = dist[pop] + 1
                    q.append(i)
    return dist[goal]
#二次元盤面の受け取り
H, W = map(int,input().split())
grid = []
for i in range(H):
    A = input()
    if 'S' in A:
        start = (i, A.index('S'))
    if 'T' in A:
        goal = (i, A.index('T'))
    grid.append(A)
#薬マスの座標を受け取り
N = int(input())
medicine = []
energy = []
for i in range(N):
    r, c, e = map(int, input().split())
    r -= 1
    c -= 1
    medicine.append((r, c))
    energy.append(e)
#スタートとゴールに薬が無い場合はエネルギー0の薬があることにする
（スタートにない時点でNoを出力してexit()でいいが気づいていない顔）
if start not in medicine:
    medicine.append(start)
    energy.append(0)
if goal not in medicine:
    medicine.append(goal)
    energy.append(0)
#薬マス間のグラフをつくる
G = [[] for i in range(len(medicine))]
for i in range(len(medicine)):
    grid_start = medicine[i]
    en = energy[i]
    dist = BFS_2Dgrid(grid,grid_start)
    for j in range(len(medicine)):
        if i == j:
            continue
        else:
            x, y = medicine[j]
            if 0 <= dist[x][y] and dist[x][y] <= en:
                G[i].append(j)
#グラフ上でスタートゴールがつながっているか判定
s = medicine.index(start)
g = medicine.index(goal)
if BFS_1Dgraph(G,s,g) >= 0:
    print('Yes')
else:
    print('No')

提出リンク：提出 #52292164 - トヨタ自動車プログラミングコンテスト2024#4（AtCoder Beginner Contest 348）