Python初心者のためのABC203

A問題
B問題
C問題
E問題

A問題

if文を使って条件を素直に実装すると良いです．

a,b,c = map(int,input().split())

if a==b:
  print(c)
elif b==c:
  print(a)
elif c==a:
  print(b)
else:
  print(0)

ソートを使って比較回数を減らすのも良いでしょう．

l = list(map(int,input().split()))
l.sort()

if l[1]==l[0]:
  print(l[2])
elif l[2]==l[1]:
  print(l[0])
else:
  print(0)

さらに簡潔に書くことも可能です．

a,b,c = map(int,input().split())
print(0 if (a-b)*(b-c)*(c-a) else a^b^c)

^はXOR演算子と呼ばれるもので詳細は割愛しますが，上記の回答は

結合法則 $(x\wedge y)\wedge z=x\wedge (y\wedge z)$ が成り立つ
$x\wedge x=0$

という性質を利用しています．

B問題

atcoder.jp

まず問題を見たときに考えるべきは愚直解（全探索，シミュレーション，等）が可能であるか？ということです．愚直解のメリットは，難しい思考やテクニックは何も必要とせず素直に実装すれば良い点です．一方のデメリットは，計算量が大幅にかかってしまうです．逆に言えば，計算量の課題さえクリアすれば愚直解が可能なので，計算量がどれくらいになるか見積もってみましょう．

今回の場合，愚直解はすべての $(i,j)$ の組について考える全探索になりますが，高々 $NK\le 9\times 9=81$ 通りしかないので計算量的には余裕です．よって，全探索で実装しましょう．問題はi0jという数値の表し方です．各桁を文字列に変換して結合した後に全体を整数に変換しても良いですが， $i0j = 100\times i + j$ を使うのが簡単でしょう．

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

ans = 0

for i in range(1,N+1):
  for j in range(1,K+1):
    ans += 100 * i + j
    
print(ans)

公式解説にある通り，数学の力を借りれば $O(1)$ で解くことも可能ですが， $O(1)$ の解法を思いついて実装するまでの時間よりも，難しい思考なしで全探索を実行する方が早いと思われます．

C問題

atcoder.jp

Step 1：愚直解の構成

まずは愚直解（シミュレーション）が可能であるかを考えましょう．太郎君は1円あたり1村進むので，すべての友達からお金を回収できるような場合に最大で $K+\sum_i B_i \le 10^9 + 10^9 \times (2\times 10^5) \simeq 10^{14}$ 村進みます．よって，シミュレーションは計算量的に不可能です．

Step 2：愚直解の改善

このシミュレーションに「無駄な計算」はないか考えてみましょう．例えば，

太郎君は現在 $M$ 円持っていて，次の友達がいる村までの距離は $D$ である

という状況では，次の友達がいる村にたどり着けないなら $M$ 回，たどり着けるなら $D$ 回のシミュレーションを行うことになります．しかし，よくよく考えてみると，次の友達がいる村にたどり着けるかどうかは $M$ と $D$ を比較すれば $O(1)$ で判定できます．この判定を行うことで，次に $M$ 回または $D$ 回のシミュレーションを行うことが分かっていれば，それらを1回のシミュレーションにまとめてしまえば良いです．

この方法により，シミュレーションの計算回数を削減することができます．具体的には，太郎君は見かけ上だと友達がいる村だけを飛び飛びで移動する（もしくは，友達がいる村にたどり着けず友達がいない村でシミュレーションが終了する）ことになるので，計算回数は $O(N)$ になります．

シミュレーション自体の計算量は $O(N)$ ですが，友達がいる村を番号が小さい順に探索するために $O(N\log N)$ の計算量をかけてソートする必要があります．

N,K = map(int,input().split())
friend = [list(map(int,input().split())) for _ in range(N)]

friend.sort()

X = 0  # 太郎君が現在いる村の番号
M = K  # 太郎君の所持金

for a,b in friend:
  
  # 太郎君が次の友達がいる村にたどり着くことができるか
  if a - X <= M:
    M += -(a - X) + b  # (a-X)村進んで(a-X)円失うが友達からb円もらえる
    X = a
    
  else:
    X += M
    print(X)
    exit()

# 最後の友達がいる村にたどり着いた後も太郎君は進み続ける
X += M
print(X)

問題を見たときに最初に愚直解を考える習慣をつけておくと，想定解が愚直解の場合に速やかに実装に取り組めるだけでなく，この問題のように想定解を構成するための原型となる解法を手に入れられる場合があります．

E問題

atcoder.jp

※難易度は青diffと初心者向きではないですが，解法を思いつきさえずれば実装は比較的簡単なので解説を書きました．

Step 1：愚直解の構成

まずは愚直解（シミュレーション）を考えます． $i$ が小さい順にすべての $j$ についてシミュレーションを行うので，計算回数は $O(N^2)$ 回です． $N\le 10^9$ なので到底間に合いません．

Step 2：愚直解の改善

C問題と同様，このシミュレーションに「無駄な計算」はないか考えてみましょう． $i+1$ 行目に黒のポーンが1つも置かれていないとします．この場合， $i$ 行目で白のポーンが置かれている可能性がある $j$ の集合を $S_i$ とすると， $S_i = S_{i+1}$ が成り立ちます．よって，

① 黒のポーンが1つも置かれていない行は削除してシミュレーションを行う

ことで，残る行は高々 $M$ 個になるため計算量を $O(NM)$ に減らすことができました．

現在は各行 $i$ に対してすべての列 $j$ についてシミュレーションを行っていますが，その必要はあるのでしょうか．例えば，0行目については $S_0=\{N\}$ です．1行目への遷移を考えるとき，そもそも $(0,j)$ に白のポーンが置かれている可能性がなければ $(0,j)$ からのシミュレーションを行う意味はありません．よって， $(0,N)$ からのシミュレーションのみを行えば良いことになります．同様に考えることで，

② $i$ 行目については $S_i$ に含まれる $j$ についてのみシミュレーションを行う

ようにすれば良いです．後の議論から $|S_i|\le M+1$ が成り立つため，全体の計算量を $O(M^2)$ に減らすことができます．

現在は白のポーン視点でシミュレーションを行っており，白のポーンが遷移する可能性がある列を探索することで $S_i$ を更新しています．一方で，黒のポーン視点で $S_i$ の更新を考えるとどうなるでしょうか．簡単のため， $i+1$ 行目には $j^\ast$ 列目にしか黒のポーンが置かれていない状況を考えてみます．すると， $j\ne j^\ast$ なる $j$ については

$j\in S_i \Longrightarrow j\in S_{i+1}$
$j\notin S_i \Longrightarrow j\notin S_{i+1}$

が成り立つことが分かります．よって， $S_i$ から $S_{i+1}$ への更新を考える上で変更（追加または削除）が行われる可能性がある要素は $j^\ast$ のみということになります．1行に複数個の黒のポーンが置かれている場合も同様に考えられるため，

③ $i$ 行目については黒のポーンが置かれている列 $j$ についてのみ $S_i$ に対する追加または削除を行う

ようにすれば良いです．これにより計算量は $O(M)$ に抑えられ時間制限に間に合うようになりました．

Step 3：実装上のテクニック

以上の議論から，①と③をふまえたプログラムを実装すれば良いことになりました．まず，①を実装するにあたっては，

黒のポーンが1つ以上置かれている行のみを抽出する
各行に置かれている黒のポーンの列番号を管理する

ような機能が必要です．これらはcollections.defaultdictによって同時に実現できます．defaultdict(func)は基本的に通常のdict()と同じですが，

辞書内に存在しないキーが呼び出されたときに，そのキーに対応する要素をfuncで初期化する

という違いがあります．funcには関数が入り，よく用いられるのはint（0で初期化），list（空のリストで初期化）などです．今回は集合を管理するためにdefaultdict(set)を使います．

from collections import defaultdict

N,M = map(int,input().split())
D = defaultdict(set)

# D[i]：i行目に置かれている黒のポーンの列番号の集合
for _ in range(M):
  x,y = map(int,input().split())
  D[x].add(y)

dictおよびdefaultdictには以下のようなメソッドがあります．

# 辞書の各要素は「キー：値」
Dict = {1:2, 'A':'B', 'list':[0,0,0]}

# 辞書のキーを出力
for k in Dict.keys():
  print(k)
> 1
> a
> list

# 辞書の値を出力
for v in Dict.values():
  print(v)
> 2
> B
> [0, 0, 0]

# 辞書の要素を出力
for k,v in Dict.items():
  print(k, v)
> 1 2
> A B
> list [0, 0, 0]

今回は，黒のポーンが1つ以上置かれている行の番号を小さい順に取得したいので，sorted(D.keys())を用いれば良いです．

次に，③の実装を考えます．実際には黒のポーンが1つ以上置かれている行しか考えないため行番号は飛び飛びになる可能性がありますが，便宜上 $S_i$ から $S_{i+1}$ への遷移を考えます．その際， $S_{i+1}$ は $S_i$ をコピーしてから $i+1$ 行目に置かれている黒のポーンを処理していけば良さそうですが，集合のコピーには計算時間がかかってしまいます．そこで，すべての行で同じ集合 $S$ を使い回すようにして，各行に関する遷移では

$S$ に新たに追加される要素
$S$ から削除される要素

を考えるようにします．

from collections import defaultdict

N,M = map(int,input().split())
D = defaultdict(set)

# D[i]：i行目に置かれている黒のポーンの列番号の集合
for _ in range(M):
    x,y = map(int,input().split())
    D[x].add(y)
    
S = set()  # 白のポーンが置かれている可能性がある列番号の集合
S.add(N)

for i in sorted(D.keys()):
    new = set()  # Sに新たに追加される列番号
    rem = set()  # Sから削除される列番号
    for j in D[i]:
        if j not in S and (j-1 in S or j+1 in S):
            new.add(j)
        if j in S and j-1 not in S and j+1 not in S:
            rem.add(j)
    for j in new:
        S.add(j)
    for j in rem:
        S.remove(j)

print(len(S))

辞書のキーのソートの部分がボトルネックとなり，全体の計算量は $O(M\log M)$ です．

きゃねろぐ

とある「おひさま」の徒然日記

Python初心者のためのABC203

A問題

B問題

C問題

E問題