Zenn
💻

ABC303 D - Shift vs. CapsLock [2次元配列を使う or 多重代入]

2023/06/01に公開
AtCoder
競プロ
動的計画法
tech

素直にDP（動的計画法）で解いていけばいい問題ではあるけれど、油断するとハマってしまうので気を付けたいというお話。（自分がやらかしたのでその記録ということで...）

問題要約 (問題へのリンク)

  • "A"と"a"からなる文字列Sが与えられる。
  • 以下のの３つの操作を繰り返すことで文字列Sを入力したい。
  • Sを入力するのにかかるコストの最小値は？

【操作】

  1. CapsLock がon[off] なら"A" ["a"]を打ち込む。コストはX。
  2. CapsLock がon[off] なら"a" ["A"]を打ち込む。コストはY。
  3. CapsLock がon[off] ならCapsLockをoff [on]にする。コストはZ。

制約

  • 1\leq |S| \leq 3 \times 10^{5}
  • 1\leq X, Y, Z \leq 10^{9}

解法

公式解説のページはこちらです。

Sのi文字目をS_{i}で表すことにします(1\leq i \leq |S| )。S_{i} まで打ち込んだとき、S_{i+1}の打ち込み方は次のパターンに限られます。

S_{i+1}= "A" S_{i+1}= "a"
CapsLockがOnになっている 操作1 or 操作3→操作2 操作2 or 操作3→操作1
CapsLockがOffになっている 操作2 or 操作3→操作1 操作1 or 操作3→操作2

これより、S_{i}まで打ち込んだ後にCapsLockがonとなっている場合とoffとなっている場合を考え、それぞれの状態に達するまでの最小コストを求めておけば、S_{i+1}まで打ち込むときの最小コストを求められると分かります。

そこで、

  • dp[i][0]: S_{i} まで打ち込んだ後にCapsLockがonとなっているようにするための最小コスト。
  • dp[i][1]: S_{i} まで打ち込んだ後にCapsLockがoffとなっているようにするための最小コスト。

として、上の表のような遷移をさせていくDPで解きます。

dpの生成と初期化

以下、Rubyで実装します。

大きさが (|S|+1) \times 2 の2次元配列としてdpを作ります。初期値はnilとしておきます。また、dp[0][0]とdp[0][1]は、それぞれZ と0 にしておきます（|S|=0 と考えれば自然？）。

dp = Array.new(S.size+1) {Array.new(2)}
dp[0][0] = Z
dp[0][1] = 0

dpの更新

上の表に基づいて、dpを更新します。各操作でCapsLockの状態が変わるかに気を付けます。また、S[i]がS_{i+1}を表しているというindexのずれにも注意します。

(S.size).times do |i|
    if S[i] == "A"
        dp[i+1][0] = [dp[i][0]+X, dp[i][1]+Z+X].min
        dp[i+1][1] = [dp[i][0]+Z+Y, dp[i][1]+Y].min
    elsif S[i] == "a"
        dp[i+1][0] = [dp[i][0]+Y, dp[i][1]+Z+Y].min
        dp[i+1][1] = [dp[i][0]+Z+X, dp[i][1]+X].min
    end
end

出力

dp[|S|][0]とdp[|S|][1]のうち小さい方が答えとなります。dp[|S|]はdp[-1]でよいです。

puts dp[-1].min

全体のコード

通して書くと次のようになります。

X, Y, Z = gets.split.map(&:to_i)
S = gets.chomp

dp = Array.new(S.size+1) {Array.new(2)}
dp[0][0] = Z
dp[0][1] = 0

(S.size).times do |i|
    if S[i] == "A"
        dp[i+1][0] = [dp[i][0]+X, dp[i][1]+Z+X].min
        dp[i+1][1] = [dp[i][0]+Z+Y, dp[i][1]+Y].min
    elsif S[i] == "a"
        dp[i+1][0] = [dp[i][0]+Y, dp[i][1]+Z+Y].min
        dp[i+1][1] = [dp[i][0]+Z+X, dp[i][1]+X].min
    end
end

puts dp[-1].min

別の実装

さて、動的計画法で解けたわけですが、今回はdp[i+1]を得るためにdp[i]の情報しか用いません。それを踏まえて次のような実装を考えます。

  1. はじめ、t=Z、u=0とする。
  2. Sの文字を先頭から1文字ずつ見ていき、
    • 見た文字が"A"ならばtをt+Xとu+Z+Xのうち小さい方に、uをt+Z+Yとu+Yのうち小さい方にそれぞれ置き換える。
    • 見た文字が"a"ならばtをt+Yとu+Z+Yのうち小さい方に、uをt+Z+Xとu+Xのうち小さい方にそれぞれ置き換える。

t=dp[i][0]、u=dp[i][1]としているわけです。1はdp[0][0]=Z、dp[0][1]=0とした部分に相当します。2はdpの更新に相当します。ただし、更新するごとに更新前の情報を捨てています。

失敗例

上記の方針を基に、次のように実装してみます。するとWAが出ます。（こんな感じで私はやらかしました。）

X, Y, Z = gets.split.map(&:to_i)
S = gets.chomp.split('')

t = Z
u = 0

S.each do |s|
    if s == "A"
        t = [t+X, u+Z+X].min
        u = [t+Z+Y, u+Y].min
    elsif s == "a"
        t = [t+Y, u+Z+Y].min
        u = [t+Z+X, u+X].min
    end
end

puts [t, u].min

問題はtとuの更新の部分です。この実装だと、先にtを更新し、更新されたtを用いてuを更新しているため間違いとなります。

失敗例の改善

この問題を回避するには多重代入を利用すればよいです。そうすることで、tとuを同時に更新でき、問題を解消できます。

X, Y, Z = gets.split.map(&:to_i)
S = gets.chomp.split('')

t = Z
u = 0

S.each do |s|
    if s == "A"
        t, u = [t+X, u+Z+X].min, [t+Z+Y, u+Y].min
    elsif s == "a"
        t, u = [t+Y, u+Z+Y].min, [t+Z+X, u+X].min
    end
end

puts [t, u].min

まとめ

ABC303 D - Shift vs. CapsLockの解法を考えてきました。素直に2次元配列を使ったDPで解けば問題ありません。しかし、2次元配列を使わずに解こうとすると、値の更新の仕方を気を付けなければやらかすという話です。

（ちなみに、「別の実装」の方がスピードは速いです。この問題ではどちらのコードでも余裕で制限時間に間に合いますが...）

教訓

  • DPは素直に多次元配列を使おう
  • 複数の変数を同時に更新したいときは多重代入を効果的に使おう

Discussion