ABC311をPythonで解いてみたよ。(A～E問題)

AtCoder Beginners Contest 311 (ABC311) をPythonで解きました。
見やすいコードを書く練習も兼ねてます。

A - First ABC

問題ページ

考察

文字 'A' について、文字列Sの左から順番に見ていって、はじめて 'A' が出てきたインデックスを記録します。
文字 'B' と 'C' についても同じようにインデックスを記録します。
この3つのインデックスの中で、最も大きな値が答えです。
左から見ていってはじめて出てくる文字を探すのは、Pythonだとfind()メソッドやindex()メソッドがあって、これをつかうとすっきり書けそうです。
(Pythonは0-indexだけど、答えを出力するときは1-indexに直さないといけないのに気をつけて！)

コード

N = int(input())
S = input()

ans = max(S.find('A'), S.find('B'), S.find('C')) + 1
print(ans)

B - Vacation Together

問題ページ

考察

$T[i]$: $i$ 日目がヒマなら'o', そうじゃないなら'x'
になるリスト $T$ を用意します。
これは文字列 $S$ を受け取るごとに、 $S$ で'x'があるインデックスはすべて $T$ でも 'x' にすることでつくれます。

リスト $T$ の中で連続する'o'の数の最大値を求めるパートは、「ランレングス圧縮」と調べると似たような問題や、それを解くコードがたくさんネットに落ちています。
下のコードでは、左から順に文字を見ていって、「今まで見た文字列について、末尾に'o'がいくつ連続してあるか」を表す変数 prev と、答えとなる変数 ans を適宜いじりながら進めています。
(自作ライブラリとして、ランレングス圧縮をする関数をあらかじめ用意しておくといいかもです...!!)

コード

N, D = map(int, input().split())

T = ['o'] * D  # T[i]: i日目が全員ヒマなら'o', そうじゃないなら'x'
for _ in range(N):
    S = input()
    for i in range(D):
        if S[i] == 'x':
            T[i] = 'x'

prev = 0  
ans = 0
for t in T:
    if t == 'o':
        prev += 1
        ans = max(ans, prev)
    else:
        prev = 0

print(ans)

C - Find it!

問題ページ

考察

適当に点1からスタートして、順番に点を通っていきます。
するとどこかのタイミングで、「この点、通ったことあるよね...？？」になります。( $N$ 頂点 $N$ 辺の有向グラフだと必ずどこかでこうなります。)
通ったことのある点に行くということは、そこに有効閉路(サイクル)があるということになるので、それを答えとして出力してプログラムを終了させればokです。

コード

N = int(input())
A = [int(t) - 1 for t in input().split()]

seen_list = [0]
se = {0}  # seen_list のセット
v = 0
while True:
    nv = A[v]
    # これ通ったことある点じゃん！のとき。
    if nv in se:
        idx = seen_list.index(nv)
        ans_list = seen_list[idx:]  # ans_list[0] = nv
        print(len(ans_list))
        print(*[a + 1 for a in ans_list])
        break
    # これは通ったことない点だなぁ。。のとき。
    else:
        se.add(nv)
        seen_list.append(nv)
        v = nv

D - Grid Ice Floor

問題ページ

考察

$(1,1)$ の点からはじめて、とりあえず上下左右4方向それぞれについて岩にぶつかるまですべってみます。
岩にぶつかったら、またそこから上下左右4方向それぞれについて岩にぶつかるまですべって、... と行けるところをすべてしらみつぶしに見ていく方法で解きます。
氷ですべったところをマーキングする2次元リスト $T$ を用意します。
あとは幅優先探索BFSのときと似た感じで、次に見る点を格納してるキューと、見たことのある点を記録するリストseenを用意します。岩にぶつかって止まったマスはキューに入れてseenに入れます。氷で滑ったマスは $T$ に記録しておきます。で、キューが空っぽになるまで探索し続けます。(これはたぶん下のコードを見てもらった方がはやいです。)

コード

from collections import deque

N, M = map(int, input().split())
S = [input() for _ in range(N)]

T = [[False] * M for _ in range(N)]
T[1][1] = True

seen = set()
vec = [(1, 0), (-1, 0), (0, 1), (0, -1)]
que = deque()
que.append((1, 1))
seen.add((1, 1))

while que:
    oi, oj = que.popleft()
    for vi, vj in vec:
        for k in range(1, 1000):
            ni, nj = oi + vi * k, oj + vj * k
            if S[ni][nj] == '.':  # 氷なのですべって次のマスへ。
                T[ni][nj] = True
            else:  # 岩なので次のマスへは進めない。
            # (ni,nj)は岩のマスなので、直前にいたマスを(ni,nj)にしてる。
                ni -= vi  
                nj -= vj
                if (ni, nj) not in seen:
                    seen.add((ni, nj))
                    que.append((ni, nj))
                break

ans = 0
for i in range(N):
    for j in range(M):
        if T[i][j]:
            ans += 1
print(ans)

E - Defect-free Squares

問題ページ

考察

「この点を左上隅とする正方形はいくつできる？」をグリッド上のすべての点で求めて足し算することで、答えを出したいと思います。
それぞれの点について、「どれだけ正方形を広げられるかな～」と1つ1つ調べていたら、今回は $H \times W \leq 9 \times 10^6$ なのでTLEしちゃいそうです。
なので、穴のあいた点を主軸において考えてみます。
例えば、とても大きなグリッド上で $(30,50)$ の点に穴があいていたとします。
そうすると、下のような影響があります。(手元の紙とかで実際に書いてみた方がいいかも!!)

• $(29,50)$ の点を左上隅とする正方形の数は $1$
• $(29,49)$ の点を左上隅とする正方形の数は $1$
• $(30,49)$ の点を左上隅とする正方形の数は $1$
• $(28,50)$ の点を左上隅とする正方形の数は $2$
• $(28,49)$ の点を左上隅とする正方形の数は $2$
• $(28,48)$ の点を左上隅とする正方形の数は $2$
• $(29,48)$ の点を左上隅とする正方形の数は $2$
• $(30,48)$ の点を左上隅とする正方形の数は $2$
• ... (以下省略)

これは、次の問題を解いているのと同じです。

穴のあいた点をスタートとして、あなたは左・上・左上に進めます。
正確には、あなたが点 $(i,j)$ にいるとき、一手で $(i,j-1), (i-1,j), (i-1,j-1)$ のいずれかに進めます。
グリッド上のすべての点について、到達できるまでの最短手数を求めてください。

これは幅優先探索BFSで解くことができます。・・・(A)

また、もう一つ問題があって、例えば $3 \times 3$ のグリッドのとき、 $(2,3)$ や $(3,1)$ を左上隅とする正方形は1つしかつくれません(グリッドが小さくて、それ以上広げられないです)。・・・(B)

$H \times W$ のグリッドの点それぞれについて、(A) の値と (B) の値の小さい方の値が、つくれる正方形の数になります((A)の制限がない点については、(B)の値だけでok)。

コード

from collections import deque

vec = [(-1, 0), (0, -1), (-1, -1)]

H, W, N = map(int, input().split())
Grid = [[-1] * W for _ in range(H)]
que = deque()
for _ in range(N):
    a, b = map(int, input().split())
    a -= 1
    b -= 1
    que.append((a, b))
    Grid[a][b] = 0

while que:
    oi, oj = que.popleft()
    for vi, vj in vec:
        ni, nj = oi + vi, oj + vj
        if not (0 <= ni < H and 0 <= nj < W):
            continue
        if Grid[ni][nj] != -1:
            continue
        Grid[ni][nj] = Grid[oi][oj] + 1
        que.append((ni, nj))


# 左上が点(i, j) になる正方形の、1辺の長さの最大値
def f(i, j):
    return min(H - i, W - j)


ans = 0
for i in range(H):
    for j in range(W):
        if Grid[i][j] == -1:
            ans += f(i, j)
        else:
            ans += min(Grid[i][j], f(i, j))
print(ans)