Zenn
🐊

ABC266 D - Snuke Panic (1D) Python解答例

2022/08/28に公開
Python
競技プログラミング
AtCoder
競プロ
tech

AtCoder Beginner Contest 266 D - Snuke Panic (1D)をPythonで解きます。

問題

問題文をAtCoderのページより引用します。

問題文

高橋君はすぬけ君たちを捕まえようとしています。

数直線上の座標0,1,2,3,4の5箇所に穴があり、すぬけ君たちの巣につながっています。

これからN匹のすぬけ君が穴から出てきます。i番目のすぬけ君は時刻T_iに座標X_iの穴から出てきて、大きさはA_iであることがわかっています。

高橋君は時刻0に座標0におり、数直線上を単位時間あたり1以下の速さで移動することができます。
すぬけ君が穴から出てきたのと同じ時刻に同じ座標に高橋君がいるとき、かつ、そのときに限り、高橋君はすぬけ君を捕まえることができます。
すぬけ君を捕まえるのにかかる時間は無視できます。

高橋君が適切に行動したとき、捕まえることができるすぬけ君の大きさの合計の最大値を求めてください。

制約

  • 1 \leq N \leq 10^5
  • 0 < T_1 < T_2 < \ldots < T_N \leq 10^5
  • 0 \leq X_i \leq 4
  • 1 \leq A_i \leq 10^9
  • 入力に含まれる値は全て整数である

解答例

動的計画法で解く

まず、考察を簡単にするため、高橋君の行動である「数直線状を単位時間あたり1以下の速さで移動する」を「各時刻について『左に1つ動く』『その場で待機する』『右に1つ動く』の3通りから選ぶ」と言い換えます。

高橋君がある時刻tにある座標xにいたとすると、高橋君の行動パターンは上記の3通りしかないので、高橋君は以下の3パターンのどれかを取って時刻tに座標xに至ることがわかります。

  1. 時刻t - 1に座標x - 1にいて、1秒かけて右に移動した
  2. 時刻t - 1に座標xにいて、その場で待機した
  3. 時刻t - 1に座標x + 1にいて、1秒かけて左に移動した

同時に、時刻tに座標xにいるためには、時刻t - 1での情報のみ必要であり、時刻t - 2以前の情報は必要としないこともわかります。

なんとなく動的計画法っぽく解けるような気がしてきました。

DPテーブルの定義

DPテーブルを以下のように定義します。尚、それまでに捕まえたすぬけ君の大きさの合計の最大値を「スコア」と呼ぶこととします。

dp[i][j]：時刻iに座標jに至ったときの、スコアの最大値。(0 \leq t \leq \max(T), 0 \leq j \leq 4)

初期値は以下のようにします。

dp[i][j] = \begin{cases} 0 \quad (i = 0 \land j = 0) \\ -\infty \quad (\text{other}) \end{cases}

また、2次元配列Cを次のように定義します。

C_{ij} = \begin{cases} A_k \quad (i = T_k, j = X_k, 1 \leq k \leq N) \\ 0 \quad (\text{other}) \end{cases}

DPの遷移は以下のようになります。

dp[i][j] = \max(dp[i - 1][j - 1], dp[i - 1][j], dp[i - 1][j + 1]) + C[i][j] \quad (1 \leq i \leq \max(T), 0 \leq j \leq 4)

\maxの中身の項はそれぞれ「右に移動してjに至る」「その場で待機してjに至る」「左に移動してjに至る」を表します。
ただし、j = 0のときとj = 4のときは領域外参照を防ぐために条件分けをしなければならないことに注意してください。

解は\max(dp[N])を出力すればよいです。

実装例

Pythonによる実装例を以下に示します。

d.py
from typing import *


def main():
    # 入力受け取り
    N: int = int(input())
    T, X, A = map(list, zip(*[list(map(int, input().split())) for i in range(N)]))

    # Tの最大値(シミュレートする時間の範囲を決める)
    time: int = max(T)

    # 2次元配列Cの準備
    C: List[List[int]] = [[0] * 5 for i in range(time + 1)]
    for t, x, a in zip(T, X, A):
        # 時刻tに座標xにすぬけ君が現れるとき、その大きさを記録する
        # 現れないなら0にする
        C[t][x] = a

    # DPテーブル
    # 無限小は適当に負の大きい数を充てておく
    dp: List[List[int]] = [[-(1 << 60)] * 5 for i in range(time + 1)]
    # 高橋君は時刻0に座標0にいる
    dp[0][0] = 0
    # 遷移開始
    for i in range(1, time + 1):
        for j in range(5):
            dp[i][j] = (
                max(dp[i - 1][max(0, j - 1)], dp[i - 1][j], dp[i - 1][min(4, j + 1)])
                + C[i][j]
            )

    # 解の出力
    print(max(dp[time]))


if __name__ == "__main__":
    main()

実際の提出はこちら。

GitHubで編集を提案

Discussion