ABC Atcoder プログラミング

【ABC267】PythonでA~D問題を解く!(AtCoder Beginner Contest 267)

2022年9月3日開催のAtCoder Beginner Contest 267(ABC267)の初心者向け解説です。今回はC問題、D問題ともに頻出のテーマの良問な印象でした。

過去のアーカイブはコチラから!

A - Saturday

問題

ある日、学校へ行くのに疲れてしまった高橋くんは、土曜日まであと何日あるかを知りたくなりました。
その日は平日で、曜日を英語で表すと \(\ S \) だったことが分かっています。その日より後の直近の土曜日は何日後かを求めてください。

なお、月曜日、火曜日、水曜日、木曜日、金曜日はそれぞれ英語で MondayTuesdayWednesdayThursdayFriday です。

制約と入力形式

  • S は MondayTuesdayWednesdayThursdayFriday のいずれかである

回答アプローチ

もちろん入力された文字列で、if文で条件分岐させれば良いですが、若干冗長かつコードも長くなってしまいす。コードが長くなると、コードが通らなかった場合にどこが悪いか見つけにくくなるので、なるべく短く書くのが良いです。

Pythonでは辞書型というデータ構造があるので、こういった問題で冗長さを回避するにはとても便利です。

またこういった問題では、文字列の綴り間違いでエラーとなるケースが発生しやすいので、なるべく条件分岐に必要な文字列は問題文からコピーしちゃいましょう。

S = input()
D = {"Monday":5,
     "Tuesday":4,
     "Wednesday":3,
     "Thursday":2,
     "Friday":1}
print(D[S])

B - Split?

問題

ボウリングのピンは \(\ 1 \)から\(\ 10 \)の番号が付けられており、上から見ると下図のように配置されます。

(7)(8)(9)(10)
(4)(5)(6)
(2)(3)
(1)

この表の縦での区切りを列と呼ぶことにします。
例えば、ピン\(\ 1,5 \) とピン\(\ 3,9 \)はそれぞれ同じ列に存在します。

いくつかのピンが倒れた状態のうち、特殊なものはスプリットと呼ばれます。
ピンの配置がスプリットであるとは、以下の条件が全て成り立つことを言います。

  • ピン\(\ 1 \) が倒れている。
  • ある二つの異なる列であって、次の条件を満たすものが存在する。
    • それぞれの列には、立っているピンが\(\ 1 \) 本以上存在する。
    • それらの列の間に、ピンが全て倒れている列が存在する。

具体例は入出力例を参考にしてください。

さて、あるピンの配置が長さ\(\ 10 \) の文字列\(\ S \) として与えられます。\(\ i = 1,2,\ldots,10 \) について、ピン\(\ i \) が倒れているとき\(\ S \) の\(\ i \) 文字目は 0 であり、ピン\(\ i \) が立っているとき\(\ S \)の\(\ i \) 文字目は 1 です。
\(\ S \)で表されるピンの配置がスプリットかどうか判定してください。

制約と入力形式

  • \(\ S \) は 0 と 1 からなる長さ\(\ 10 \)の文字列

回答アプローチ

各列で、全部倒れてるか否かを出します。全部で\(\ 7 \) 列で、一本でも立っている列は\(\ 1 \)、全て倒れている列は\(\ 0 \)とします。これを再度文字列に戻してしまいます。例えば、 両端だけが残っている場合は、「1000001」と表されます。

列単位で見た時スプリットとなりうるパターンは、「1000001」「100001」「10001」「1001」「101」の精々\(\ 5 \) パターンであり、これが上記の文字列に含まれるか否かを判定すれば良いだけです。

AtCoderのPyPyでは、exit()でプログラムを終了できるので、スプリットが見つかったらprint("Yes")をして、プログラムを終了させましょう。

S =list(map(int,list(input())))
C = [S[6],
     S[3],
     S[1] or S[7],
     S[0] or S[4],
     S[2] or S[8],
     S[5],
     S[9]]
column = "".join(map(str,C))
split_list = ["1"+"0"*i+"1" for i in range(1,6)]
if S[0] == 0:
  for sl in split_list:
    if sl in column:
      print("Yes")
      exit()
print("No")

C - Index × A(Continuous ver.)

問題

長さ \(\ N \) の整数列\(\ A = (A_1, A_2, A_3, \ldots , A_N )\)が与えられます。

長さ \(\ M \) の \(\ A \) の連続部分列\(\ B = (B_1, B_2, B_3, \ldots , B_N) \)に対する、\(\ \sum_{i=1}^M i \times B_i \) の最大値を求めてください。

制約と入力形式

  • \(\ 1 \leq M \leq N \leq 2 \times 10^5 \)
  • \(\ -2 \times 10^5 \leq A_i \leq 2 \times 10^5 \)
  • 入力は全て整数。

回答アプローチ

\(\ \sum_{i=1}^M i \times B_i \) が最大となる連続部分列 \(\ B \) を全探索する問題です。ただし、毎回 \(\ \sum_{i=1}^M i \times B_i \) を計算してしまうと、\(\ O(N \times M) \) となり、TLEとなります。このような長い数列上を探索する問題では、「累積和」の考え方を使うのが有効であることが多いです。

下記図は \(\ N=5 \)、 \(\ M=4 \) の場合に、部分列 \(\ (A_1, A_2, A_3,A_4) \) の \(\ \sum_{i=1}^M i \times B_i \)を青色の四角で表しています。

下記は部分列 \(\ ( A_2, A_3,A_4,A_5) \) の \(\ \sum_{i=1}^M i \times B_i \)を桃色の四角で表して、重ねています。

桃色の部分を青色の部分を使って求めるには以下のような処理が必要となります。

青色の部分から 黄色の部分( \(\ (A_1, A_2, A_3,A_4) \) の和)を引いて、緑色の部分( \(\ A_5 \times M \) )を足すことで、桃色が求めることができます。

初項(\(\ \sum_{i=1}^M i \times A_i \) ) だけに関しては愚直に計算する必要がありますが、後続に関しては\(\ O(1) \) で求められます。ただし、黄色の部分は毎回愚直に和をを求めると\(\ M \) に比例して計算量が大きくなるので、ここを見落とすと結局ボトルネックとなります。これに関しても初項だけは愚直に和をとり、後続に関しては、最もインデックスの小さい項を引き、次の項を足すという計算で\(\ O(1) \)で計算できます。

n,m = map(int,input().split())
A = list(map(int,input().split()))
 
sigma=sum([(i+1)*a for i,a in enumerate(A[:m])]) #最初だけO(M)で計算
sumB = sum(A[:m]) #累積和
result =sigma
 
for i in range(n-m):
  sigma = sigma - sumB+ m*A[i+m]
  sumB = sumB - A[i] + A[i+m] #最もインデックスの小さい項を引き、次の項を足して更新
  if sigma>result:
    result = sigma
print(result)

D - Index × A(Not Continuous ver.)

問題

長さ \(\ N \) の整数列\(\ A = (A_1, A_2, A_3, \ldots , A_N )\)が与えられます。

長さ \(\ M \) の \(\ A \) の部分列(連続じゃなくても良い)\(\ B = (B_1, B_2, B_3, \ldots , B_N) \)に対する、\(\ \sum_{i=1}^M i \times B_i \) の最大値を求めてください。

制約と入力形式

  • \(\ 1 \leq M \leq N \leq 2000 \)
  • \(\ -2 \times 10^5 \leq A_i \leq 2 \times 10^5 \)
  • 入力は全て整数。

回答アプローチ

C問題と異なり、連続じゃなくても良くなります。ただし、順番だけは元の数列を守必要があります。組み合わせ的には \(\ _N C_M \) と膨大なので、C問題のような全探索は不可能です。ここで、動的計画法を使うことを検討します。

ややメタ的ではありますが、どのような状態遷移を組み立てるか、より先に \(\ N \times M \) の二次元配列を考えると初学者などは解きやすいと思います。

dp = [[None]*N for i in range(M)]

DPは最初は問題を小さく考えて、徐々に大きく拡張していくやり方です。今回の場合は、数列 \(\ A \) の\(\ n \) 番目までだけに着目して部分列の長さが\(\ m \) だとした時の最適解は二次元配列に埋めていきます。

下記は入力例にもある、\(\ N = 4, M=2\)の例です。入力数列が\(\ (5,4,-1,8) \) です。

例えば、DPの一行目は簡単で、部分列の長さが \(\ 1 \) の時を考えます。つまり、一個しか選べません。なので数列 \(\ A \) の\(\ n \) 番目までだけに着目した時、それまでの最大値が最適解となります。上の図では、どこを着目しても\(\ 5 \) を選ぶのが最適だということがわかります。ちなみに、\(\ M = 2 \) なので、\(\ 8 \) を部分列の一つ目の要素に選べことはできませんので空白となってます。

二行目以降は、\(\ n \) 番目の要素を採用した方が最適か否かを判定します。採用する場合の値は、上図の緑の矢印のように、現在のマスから左上の要素と足し合わせた値です。採用しない場合は、左のマスの値を引き継ぎます。採用する場合、しない場合の内大きい値をマスに入れます。このような処理を再起的に行なっていきます。

また、今回は \(\ -2 \times 10^5 \leq A_i \leq 2 \times 10^5 \) と数列の要素が負の値をとり得るので、初期化はNoneで行います。(負の値を取らない場合は、\(\ -1 \) とか\(\ 0 \) で初期化した方が何かと便利です。)

n,m = map(int,input().split())
A = list(map(int,input().split()))
dp = [[None]*n for i in range(m)]
dp[0][0] = A[0]
for i,a in enumerate(A[1:-m+1]):
  dp[0][i+1] = max(a,dp[0][i])
 
for i_m in range(1,m):
  for i_n in range(i_m,n - m + i_m+1):
    if dp[i_m][i_n-1] == None:
      dp[i_m][i_n] = dp[i_m-1][i_n-1]+(i_m+1)*A[i_n]
      continue
    else:
      dp[i_m][i_n] = max(dp[i_m][i_n-1],
                       dp[i_m-1][i_n-1]+(i_m+1)*A[i_n])
print(dp[-1][-1])

-ABC, Atcoder, プログラミング
-, ,