30. Paint Fence

読了目安約2分

DP / LC 276 (Easy) — 柵を k 色で塗る場合の数。2 状態の DP。

この章の目次

問題

LeetCode 276. Paint Fence（Easy / カテゴリ: DP）

LeetCode Premium 限定です。問題の要点は次の要約の通りです。

一列に並んだ n 本の柵の柱を、k 色の絵の具で全部塗ります。ただし、同じ色を 3 本以上連続して塗ってはいけません。同じ色が 2 本続くのは許されます。条件を満たす塗り方の総数を返します。例えば n = 3, k = 2 では、2 × 2 × 2 = 8 通りのうち「3 本とも同色」の 2 通りだけが禁止で、答えは 6 です。 n は小さく（最大 50）、答えは Int に収まります。

無料で読める類題に LeetCode 70. Climbing Stairs があります。直前の結果から次を作る、同じ型の動的計画法（DP）です。

考え方

ヒント 1

i 本目までの塗り方の数を、i - 1 本目までの結果から作れないか考えます。ただし総数 1 つだけでは足りません。 i 本目を直前と同じ色にできるかどうかは、直前 2 本の状態で決まります。

ヒント 2

状態を 2 つに分けます。 same = 最後の 2 本が同色である塗り方の数、diff = 最後の 2 本が違う色である塗り方の数。次の柱を塗るときの遷移は same' = diff（3 連続禁止のため）、diff' = (same + diff) × (k - 1) です。 n = 2 の初期値は same = k、diff = k × (k - 1) で、答えは same + diff です。

Swift 実装のポイント

配列は不要です。直前の same / diff の 2 変数だけで更新できます。
(same, diff) = (diff, (same + diff) * (k - 1)) とタプル代入すると、一時変数なしで同時に更新できます。
n == 1 だけ場合分けが必要です（答えは k）。k == 1 は特別扱い不要で、3 本目以降は漸化式が自然に 0 になります。

模範解答

Swift
class Solution {
    func numWays(_ n: Int, _ k: Int) -> Int {
        if n == 1 { return k }
        var same = k
        var diff = k * (k - 1)
        for _ in 0..<(n - 2) {
            (same, diff) = (diff, (same + diff) * (k - 1))
        }
        return same + diff
    }
}

計算量: O(n)。定数個の変数を n − 2 回更新するだけです。