2次元配列を全探索する際にネストを浅くする

2次元配列を全探索する際にネストを浅くする方法を解説します。この記事では、Pythonを使います。

AtCoder Beginner Contest 383のB問題を使って解説します。まずはリンクから問題を読んでください。

次に解説を読んでください。

この解説にある解答例では、forループが6階の入れ子になっています。1つ目の加湿器の場所について高さと幅でforループ2回、2つ目の加湿器で2回、その場所が加湿されているかどうかで2回の計6回のforループです。forループのネストを浅くするにはどうすればいいでしょうか。

ここで、次の点に着目します。$S$の全ての行を一列に並べた長さ$H \times W$の文字列$S'$を考えると、ある$i(0\le i < H \times W)$番目の文字$S_i'$と元の2次元配列$S$の$h$($i$を$W$で割った商)行目、$w$($i$を$W$で割った余り)列目の文字$S_{hw}$は一致します。

forループを$S'$の要素について行うことでネストの深さを半分にできます。

サンプルコードは以下のとおりです。Pythonではdivmodを使うことで商と余りを同時に取得できます。

abc383_b.py

H, W, D = map(int, input().split())
S = [input() for i in range(H)]

answer = 0

for i in range(H * W):
    for j in range(i + 1, H * W):
        ih, iw = divmod(i, W)
        jh, jw = divmod(j, W)
        if S[ih][iw] == "#": continue
        if S[jh][jw] == "#": continue

        count = 0
        for k in range(H * W):
            kh, kw = divmod(k, W)
            if S[kh][kw] == "." and (abs(kh-ih)+abs(kw-iw)<=D or abs(kh-jh)+abs(kw-jw)<=D):
                count += 1
        answer = max(count, answer)
print(answer)