ABC360(Atcoder Beginner Contest)のA~E(A,B,C,D,E)問題をPythonで解説(復習)

A問題

ご飯 R が一番左にあれば，必ず味噌汁 M は右側に行く．
味噌汁 M が一番右にあれば，必ずご飯 R は左側に行く．

A.py

"""
<方針>
- ご飯 `R` が一番左にあれば，必ず味噌汁 `M` は右側に行く．
- 味噌汁 `M` が一番右にあれば，必ずご飯 `R` は左側に行く．
- 
"""
# 入力を受け取る．
S = input()

# ご飯が一番左 or 味噌汁が一番右 の時，
if(S[0]=="R" or S[2]=="M"):
  # Yes を出力
  print("Yes")
else:
  # No を出力
  print("No")

B問題

文字を見る間隔 w を 1<=w<|S| で順番に変化させる．
文字を見始める位置 c を 0<=c<w で順番に変化させる．
文字を取得する位置 i を位置 c から始めて， w 間隔で変化させ，一時的にリスト tmp に保存する．
リストの文字を結合させて，文字列 T と合致したら， Yes を出力して，プログラムを終了させる．
合致がしなければ，最後に No を出力する．

B.py

"""
<方針>
- 文字を見る間隔 `w` を `1<=w<|S|` で順番に変化させる．
- 文字を見始める位置 `c` を `0<=c<w` で順番に変化させる．
- 文字を取得する位置 `i` を位置 `c` から始めて， `w` 間隔で変化させ，一時的にリスト `tmp` に保存する．
- リストの文字を結合させて，文字列 `T` と合致したら， `Yes` を出力して，プログラムを終了させる．
- 合致がしなければ，最後に `No` を出力する．
"""

# 標準入力から受け取る．
S, T = input().split()

# 文字を見る間隔
for w in range(1, len(S)):
  # 文字を見始める位置
  for c in range(0, w):
    # 一時的に文字を格納する場所
    tmp = []
    # 文字を取得する位置
    i = c
    # 文字を取得できる範囲内の時，
    while i<len(S):
      # 文字をリストに追加する．
      tmp.append(S[i])
      # 文字を参照する位置を進める．
      i += w
    
    # リストの文字列とTが合致するかどうか．
    if("".join(tmp)==T):
      # 合致したら，Yesを出力して，プログラムを終了させる．
      print("Yes")
      exit()
# 一度も合致しなければ，Noを出力して終了させる．
print("No")

C問題

方針

箱の数と荷物の数が一緒なので，最終的に，箱一つに荷物がちょうど一つ入っていれば良い．
荷物が重複している分だけ必ず，ちょうどその数の空き箱が存在する．
動かした荷物の重さ W[i] 分のコストがかかるので，重複している箱から，荷物をどかす時は，任意の空き箱に入れれば良い．
従って，荷物の重複がある箱が存在しているとき，その箱から最も重い荷物 W[i] 以外を退かせば良い．

前提

C問題あたりで，TLEになる人は，制約条件を見る癖をつけよう．
AとB問題では，基本的に制約条件を見ずにやっても解ける．
しかし，C問題以降では，制約条件を見ないと必ずTLEすると思っても良い．
詳しい話は私の352回の記事のC問題の解説に記したので，是非参照してほしい．

C.py

"""
<方針>
- 箱の数と荷物の数が一緒なので，最終的に，箱一つに荷物がちょうど一つ入っていれば良い．
- 荷物が重複している分だけ必ず，ちょうどその数の空き箱が存在する．
- 動かした荷物の重さ `W[i]` 分のコストがかかるので，重複している箱から，荷物をどかす時は，任意の空き箱に入れれば良い．
- 従って，荷物の重複がある箱が存在しているとき，その箱から最も重い荷物 `W[i]` 以外を退かせば良い．
- 
"""
# 入力受け取り．
N = int(input())
A = list(map(int, input().split()))
W = list(map(int, input().split()))

# 箱を表現するリスト．
boxes = [[] for i in range(N)]

# 最初に置いてある荷物の重さ W[i] を格納する．
for i in range(N):
  # 荷物の初期位置
  a = A[i]-1
  # 荷物の重さ
  w = W[i]
  
  # 箱の中に荷物を入れる．
  boxes[a].append(w)

# 箱の中の荷物を軽い順番でソートする．
for i in range(N):
  boxes[i].sort()

# 移動する荷物の総重量
ans = 0
# 箱の中に二つ以上の荷物があった時，一番重い荷物以外を総重量に加える．
for i in range(N):
  # 箱を選択する．
  box = boxes[i]
  # 箱の中に荷物が二つ以上ある時，
  if(len(box)>=2):
    # 一番重い荷物以外を総重量に加える．
    ans += sum(box[:-1])
    
# 総重量を出力する．
print(ans)

D問題

負の方向に移動する🐜が初期位置で任意の区間に何匹いるかを，定数時間で取れるようにしておく．
正の方向に移動する🐜を一匹ずつみて， 2*T+1 先までの区間で何匹の負の方向に進む🐜が存在するかを見る．

D.py

"""
<方針>
- 負の方向に移動する🐜が*初期位置で*任意の区間に何匹いるかを，定数時間で取れるようにしておく．
- 正の方向に移動する🐜を一匹ずつみて， `2*T+1` 先までの区間で何匹の負の方向に進む🐜が存在するかを見る．
"""
"""めぐる式二分探索"""
def bs(a, sep, rightInclude=True, rev=False):
  left = -1
  right = len(a)
  
  while (right - left > 1):
    mid = left + ((right - left) // 2)
    
    # 大>小 モード
    if(rev):
      if(a[mid] < sep): right = mid
      elif(a[mid] == sep and rightInclude): right = mid
      else: left = mid
    # 小<大 モード
    else:
      if(a[mid] > sep): right = mid
      elif(a[mid] == sep and rightInclude): right = mid
      else: left = mid
  
  # sep境界値とした左右
  return left, right

# 入力を受け取る．
N, T = map(int, input().split())
S = list(input())
X = list(map(int, input().split()))

# Xが昇順になるようにソート
X, S = zip(*sorted(zip(X, S)))

# 負の方向に動く🐜の初期位置を取得．
ls = [0]*N
for i in range(N):
  if(S[i]=="0"):
    ls[i] += 1

# 右側からの負の方向に動く🐜の初期位置の数の累積和
R = [0]*N
# 最右index
d = N-1 
# 累積和の開始
R[d] = ls[d]
for i in range(N-1):
  R[d - (i+1)] = R[d - i] + ls[d - (i+1)]
# 右端に0匹を追加する．
R.append(0)

# 正の方向に動く🐜を一匹ずつ見る．
ans = 0
for i in range(N):
  # 正の方向に動く🐜の時，
  if(S[i]=="1"):
    # 長さ2*T+1の区間に負の方向に動く🐜が何匹いるかを取得する．
    _, b = bs(X, X[i]+2*T+1)
    ans += R[i]-R[b]

# 出力
print(ans)

E問題 | 線形時間 `O(K)`

愚直に， K の線形時間(O(K))で処理をしています．

E.py

"""
<方針>
- 公式解説と一緒です．
"""
N, K = map(int, input().split())
mod = 998244353

# K回目の操作終了後に黒玉が一番左にある確率
dp = [0 for i in range(K+1)]

# 操作していない状態(初期状態)は黒玉は必ず最左にある．
dp[0] = 1

# 黒いボールが最も左にある状態から 1 回操作して、黒いボールが最も左以外のどこかに移動する確率
p = (2*(N-1)*pow(N**2, -1, mod))%mod
# 黒いボールが最も左以外のどこかにある状態から 1 回操作して、黒いボールが最も左に移動する確率
q = (2*pow(N**2, -1, mod))%mod

# K回の操作を繰り返す．
for i in range(K):
  dp[i+1] = (((1-p)*dp[i])%mod + (q*(1-dp[i]))%mod)%mod

# 期待値を求める．  
ans = (dp[K]*1 + (1-dp[K])*(N*(N+1)//2 - 1)*pow(max(N-1, 1), -1, mod))%mod

print(ans)

E問題 | 対数時間 `O(logK)`

一般項を求めることで， K の対数時間(O(logK))で処理ができます．

E(対数時間O(logK)).py

"""
<方針>
- 公式解説と一緒です．
"""
N, K = map(int, input().split())
mod = 998244353

# 黒いボールが最も左にある状態から 1 回操作して、黒いボールが最も左以外のどこかに移動する確率
p = (2*(N-1)*pow(N**2, -1, mod))%mod
# 黒いボールが最も左以外のどこかにある状態から 1 回操作して、黒いボールが最も左に移動する確率
q = (2*pow(N**2, -1, mod))%mod
# aは次の説明で用いられる特殊解
a = (q*pow(p+q, -1, mod))%mod

# K回の操作を繰り返すが，
# dp[K] = a + (1-p-q)^K(dp[0]-a)
# と，一般項が求まるので，
dpK = a + pow(1-p-q, K, mod)*(1-a)

# 期待値を求める．  
ans = (dpK*1 + (1-dpK)*(N*(N+1)//2 - 1)*pow(max(N-1, 1), -1, mod))%mod

print(ans)

補足

関係するリンク(参考文献など)

筆者について

Atcoderアカウント
今回の成績(oooox--) | B問題で AC 回答が WA になる不具合のせいで，まだ成績が出てません．パフォーマンス予想拡張機能では， 900 くらいでした．．．
解説で示したA問題の提出
解説で示したB問題の提出
解説で示したC問題の提出
解説で示したD問題の提出
解説で示したE問題(線形時間)の提出
解説で示したE問題(対数時間)の提出

その他

間違いを含んでいる可能性があります．
方針と言いつつ，方針ではない感想などを書いている可能性があります．
A問題から解説がだんだん省略されます．
コードに書かれている解説の文言と，その上に書いてある解説の文言を変えている場合があります．

最後に一言

とにかく継続することが重要です．
十で神童　十五で才子　二十過ぎれば只の人