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

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

A - Scoreboard

問題

$i=1,2,\cdots N$ について、$i$ 回目の試合ではチーム高橋が $X_i$ 点を、チーム青木が $Y_i$ 点を獲得しました。
得点の合計が多い方の勝ちです。どっちが勝ったでしょう？

考察

for文をつかって、チーム高橋とチーム青木のそれぞれの得点を更新していきます。

コード

N = int(input())
X, Y = 0, 0
for _ in range(N):
    x, y = map(int, input().split())
    X += x
    Y += y
if X > Y:
    print("Takahashi")
elif X < Y:
    print("Aoki")
else:
    print("Draw")

別解

sum関数をつかった解法

内包表記とsum関数をつかって下のように書くこともできます。
この問題ではそんなに恩恵を感じないですが、内包表記もsum関数もおぼえておいて損はないものなので、この機会におぼえておきましょう。

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

X, Y = (sum(x), sum(y) for x, y in XY)
if X > Y: print("Takahashi")
elif X < Y: print("Aoki")
else: print("Draw")

B - Extended ABC

問題

文字列 $S$ のすべての文字がAのとき、 $S$ を拡張A文字列といいます。
拡張B文字列、拡張C文字列についても同様に定義します。
ある拡張A文字列 $S_A$ , 拡張B文字列 $S_B$ , 拡張C文字列 $S_C$ があって、 $S_A, S_B, S_C$ がこの順に並んでいるとき、その文字列を拡張ABC文字列といいます。
入力で与えられる文字列 $S$ は拡張ABC文字列ですか？

考察

拡張ABC文字列は、昇順にソートしても形が変わりません。
これを利用するとACできます。

コード

S = list(input())

T = sorted(S)
if S == T:
    print("Yes")
else:
    print("No")

別解

collections.Counterをつかった解法

collections.Counterをつかうと、文字列 $S$ の中にあるそれぞれの文字の個数を知ることができます。

これをつかって、 $S_A, S_B, S_C$ をつくり、$S_A+S_B+S_C$ が $S$ と一致するかどうかを見ればACできます。

from collections import Counter

S = input()

C = Counter(S)
T = 'A' * C['A'] + 'B' * C['B'] + 'C' * C['C']
if S == T:
    print("Yes")
else:
    print("No")

特定の文字列が存在しないことを利用する解法

拡張ABC文字列には、"BA", "CA" "CB" の文字列は存在しません。

S = input()

for t in ["BA", "CA", "CB"]:
    if t in S:
        print("No")
        exit()
print("Yes")

C - Lining Up 2

問題

数列 $A=(A_1, A_2, \cdots , A_N)$ は、次のことを表しています。

• $A_i=-1$ のとき、人 $i$ は列の先頭。
• $A_i \neq -1$ のとき、人 $i$ は人 $A_i$ のすぐ後ろにいる。

数列 $A$ から、列を再現してください。

考察

列を先頭から再現していきます。
列の先頭の人は、 $A_i=-1$ となっている $i$ を見つけるだけでいいです。
問題はその先で、「列の $i$ 番目の人が分かっているとき、列の $i+1$ 番目はだれでしょうか？」の問題に高速に答えたいです。

でも、数列 $A$ からわかるのは、「列の $i+1$ 番目の人が分かっているとき、列の $i$ 番目はだれでしょうか？」の問題です。
なので、これを逆にしたのを辞書に記録します。
コードにするとこんな感じ。

"""数列Aから、逆引きの辞書をつくる"""
dic = dict()
for i in range(N):
    dic[A[i]] = i

これで、列の $i$ 番目の人が分かれば列の $i+1$ 番目の人が分かるようになり、問題を解くことができます。

コード

N = int(input())
A = list(map(lambda x: int(x) - 1, input().split()))

dic = dict()
for i in range(N):
    dic[A[i]] = i
ans_lst = [dic[-2]]
for _ in range(N - 1):
    tail = ans_lst[-1]
    ans_lst.append(dic[tail])
for i in range(N):
    ans_lst[i] += 1
print(*ans_lst)

D - Cheating Gomoku Narabe

問題

$H$ 行 $W$ 列のグリッドがあり、各マスは o, x, . のいずれかです。
以下の操作を $0$ 回以上行います。

• . が書かれているマスを $1$ つ選び、そのマスを o に書き換える。

縦方向または横方向に連続した $K$ 個のマスであってそれらに書かれた文字がすべて o であるようなものが存在するようにすることができますか？できるならば、操作回数の最小値はいくつですか？

考察1: 次元を落とす

競プロの考察あるあるのうちの1つに、「2次元の問題は、まず1次元ならどう解くか考える」というのがあります。
この問題だと、 $H=1$ でどう解くかを考えてみます。

考察2: 1次元の問題を解く

連続した $K$ マスの中に x が入っていると、操作をしてもすべてのマスを o にはできません。
また、連続した $K$ マスの中に x がないとき、その中に含まれる . の数が操作回数になります。

つまり、 $1$ ～ $K$ 文字目の中で

• x があるかどうか
• . の個数はいくつか

の2つが知りたいです。

$2$ ～ $K+1$ 文字目でもどうなのか知りたいし、$3$ ～ $K+2$ 文字目でもどうなのか知りたいし。最後の $W$ 文字目までずっとこの2つが知りたいです。

このパターンはしゃくとり法がつかえます。

いろいろな書き方があると思いますが、今回はcollections.Counterをつかって、o, x, . それぞれの個数を管理することにします。コードにするとこんな感じ。

"""考察2までのコード"""
from collections import Counter

INF = 1 << 63

# (小問題に分けて解くときは、関数化しておくと実装がラクになることが多いです。)
def row_solve(row):
    w = len(row)
    if w < K:
        return INF
    result = INF
    c = Counter(row[:K])
    for i in range(w - K + 1):
        if c['x'] == 0:
            result = min(result, c['.'])
        if i < w - K:
            c[row[i]] -= 1
            c[row[i + K]] += 1
    return result

考察3: 2次元の問題に戻す

縦方向に $K$ マス選ぶことはいったんなかったことにして、横方向に $K$ マス選ぶのだけを考えればいいという条件で $H>1$ になったときのことを考えます。
1行だけの問題は解けているので、各行に対して上で書いたrow_solve関数をつかえばokです。

# 横方向だけを考えて、H>1のときにも解けるようにした関数
def solve(T):
    w = len(T[0])
    if w < K:
        return INF
    result = INF
    for row in T:
        result=min(result, row_solve(row))
    return result

あとは縦方向についても考えればゴールです。
$H$ 行 $W$ 列のグリッドを $90$ 度回転させて $W$ 行 $H$ 列のグリッドにすれば、回転前のグリッドに対して縦方向を考えるのと、回転後のグリッドに対して横方向を考えるのは一緒です。
横方向だけの問題は上で書いたsolve関数がつかえるので、$90$ 度回転させたグリッドに対してsolve関数をつかえば縦方向について考えたことになります。

下のコードでは、 $90$ 度回転させたグリッドを考える代わりに、コードを簡潔にするために $S$ を転置させた行列(反時計回りに $90$ 度回転させて左右反転させたもの)をつかっています。
もちろん(反)時計回りに $90$ 度回転させただけのものをつかってもokです。

コード

from collections import Counter

INF = 1 << 63


def row_solve(row):
    w = len(row)
    result = INF
    c = Counter(row[:K])
    for i in range(w - K + 1):
        if c['x'] == 0:
            result = min(result, c['.'])
        if i < w - K:
            c[row[i]] -= 1
            c[row[i + K]] += 1
    return result


def solve(T):
    w = len(T[0])
    if w < K:
        return INF
    result = INF
    for row in T:
        result=min(result, row_solve(row))
    return result


H, W, K = map(int, input().split())
S = [input() for _ in range(H)]
S2 = [list(a) for a in zip(*S)]  # Sの転置行列 
ans = min(solve(S), solve(S2))
if ans == INF:
    print(-1)
else:
    print(ans)

E - Bad Juice

問題

$N$ 本のジュースのうち $1$ 本だけが腐っており、飲むとお腹を壊します。
どれが腐ったジュースかを見抜きたいです。友人を 必要な最小の数 だけ呼んで、$N$ 本のジュースを好きなように飲ませるとき、その人数と、どのようにして腐ったジュースを特定するのかを教えてください。

考察1: 呼ぶ友人の数を求める

どのようにジュースを飲ませるかは後で考えるとして、いったん呼ぶ友人の数 $M$ について考えます
$M=1$ のとき、$1$ 人の友人がお腹を壊すか壊さないかの $2$ 通りあり、ジュース $2$ 本まで判別できます。
$M=2$ のとき、$2$ 人の友人のお腹を壊すかどうかのパターンは $2^2=4$ 通りあり、ジュース $4$ 本まで判別できます。
$M=3$ のとき、$3$ 人の友人のお腹を壊すかどうかのパターンは $2^3=8$ 通りあり、ジュース $8$ 本まで判別できます。

呼ぶ友人の数が $M$ のとき、ジュース $2^M$ 本までなら判別できることになります。

よって、 $N \leq 2^M$ を満たす最小の $M$ が呼ぶ友人の数です。

ここまでのコードはこんな感じ。

"""考察1までのコード"""
N = int(input())

M = len(bin(N - 1)[2:])
print(M)

考察2: 友人にどのジュースを飲ませるか決める

$M$ 個の数字があって、それぞれ $0$ か $1$ (お腹を壊すか壊さないか)で、全体として $2^M$ 通りのパターンを表せる。これは2進数と同じです。

友人 $1$ ～ $M$ をここでは便宜上、友人 $0$ ～ $M-1$ とします。
ジュースも同じく、ジュース $0$ ～ $N-1$ とします。
友人 $i$ には、ジュース $0$ ～ $N-1$ のうち、その番号を2進表記して下から $i$ 桁目が $1$ になるものを飲ませることにします。
たとえば、友人 $2$ は、ジュース$4,5,6,7,12,13,14,15,20,21,\cdots$ を飲ませることにします。
こうすれば、友人 $i$ がお腹を壊したとき、腐っているジュースの番号を2進表記すると下から $i$ 桁目は $1$ になっていることが分かります。
逆にお腹を壊していなければ、下から $i$ 桁目は $0$ だと分かります。

たとえばジュース $37$ が腐っていたとき、$37=(100101)_2$ なので、友人 $0,2,5$ がお腹を壊します。

ここまでの議論はすべて0-indexedで進めてきたので、出力の際は1-indexedに戻すのを忘れないようにしましょう。

コード

N = int(input())

M = len(bin(N - 1)[2:])
print(M)

for i in range(M):
    lst = []
    for j in range(N):
        if (j >> i) & 1:
            lst.append(j + 1)
    print(len(lst), *lst)

S = input()
ans = int(S[::-1], 2) + 1
print(ans)