概要

\(N\) 枚の板を \(K\) 色で塗り分ける問題で、隣り合う板は異なる色でなければならないという制約のもとでコストの総和を最小化します。典型的な区間 DP（列に沿った DP）ですが、素朴な遷移では間に合わないため「最小値・2番目の最小値」を使って遷移を高速化します。

考察

DP の定義

左から順に決めていくのが自然です。板 \(i\) を色 \(j\) で塗ったときの、板 \(1\) から板 \(i\) までの最小コストを

\[\mathrm{dp}[i][j]\]

と定義します。すると遷移は

\[\mathrm{dp}[i][j] = C_{i,j} + \min_{j' \neq j} \mathrm{dp}[i-1][j']\]

となります（板 \(i-1\) の色 \(j'\) は \(j\) 以外なら何でもよい）。初期値は \(\mathrm{dp}[1][j] = C_{1,j}\)、答えは \(\min_j \mathrm{dp}[N][j]\) です。

素朴な実装はなぜダメか

上の式をそのまま実装すると、各 \(i\) について各 \(j\) ごとに \(\min_{j'\neq j}\) を \(O(K)\) かけて計算するので、全体で \(O(N K^2)\) になります。

制約は \(N \times K \leq 10^6\) なので、たとえば \(N = 1, K = 10^6\) や \(N = 1000, K = 1000\) が起こり得ます。\(N=1000, K=1000\) のとき \(NK^2 = 10^9\) となり、確実に TLE です。

高速化のアイデア

\(\min_{j' \neq j} \mathrm{dp}[i-1][j']\) は「\(j\) を除いた最小値」です。ここで重要な観察は：

• 配列全体の最小値を与えるインデックスを \(i_1\)、その最小値を \(m_1\)、\(i_1\) を除いた最小値（＝2番目に小さい値）を \(m_2\) とすると、
  • \(j \neq i_1\) のときは、最小値の位置 \(i_1\) が除外されていないので \(\min_{j'\neq j} = m_1\)
  • \(j = i_1\) のときだけ、\(\min_{j'\neq j} = m_2\)

つまり「\(j\) を除いた最小値」は \(j\) によって 2 通りしかありません。よって

\[\mathrm{dp}[i][j] = \begin{cases} C_{i,j} + m_2 & (j = i_1) \\ C_{i,j} + m_1 & (j \neq i_1)\end{cases}\]

と、\(O(K)\) で 1 段分の遷移が計算できます。

具体例

\(K=4\) で \(\mathrm{dp}[i-1] = [7, 3, 9, 5]\) の場合、\(m_1 = 3\)（位置 \(i_1 = 2\)、1-indexed）、\(m_2 = 5\) です。 板 \(i\) を色 \(2\) 以外で塗るなら前の板は色 \(2\) にしてよいのでコスト \(+3\)、色 \(2\) で塗るときだけ \(+5\) となります。

アルゴリズム

1. \(\mathrm{dp} = (C_{1,1}, \dots, C_{1,K})\) で初期化する。
2. \(i = 2, 3, \dots, N\) について以下を行う。
   • \(\mathrm{dp}\) の最小値 \(m_1\) とその位置 \(i_1\)、および \(i_1\) を除いた最小値 \(m_2\) を求める。
   • 新しい配列 \(\mathrm{ndp}\) を \(\mathrm{ndp}[j] = C_{i,j} + m_1\)（ただし \(\mathrm{ndp}[i_1] = C_{i,i_1} + m_2\)）として作る。
   • \(\mathrm{dp} \leftarrow \mathrm{ndp}\)。
3. \(\min_j \mathrm{dp}[j]\) を出力する。

なお \(K \geq 2\) が保証されているので、\(m_2\) は必ず存在し（\(K \ge 2\) なら \(i_1\) 以外の要素が 1 個以上ある）、また塗り分けは常に可能です。

計算量

• 時間計算量: \(O(NK)\)（\(NK \leq 10^6\) なので十分高速）
• 空間計算量: \(O(K)\)（\(\mathrm{dp}\) 配列は 1 行分だけ持てばよい。入力の読み込み分を含めれば \(O(NK)\)）

実装のポイント

• \(N = 1\) の場合：遷移が一度も起きないので、単に \(\min_j C_{1,j}\) を出力します（ループの初期化だけで処理できるようにしておくとバグりにくい）。上のコードでは明示的に分岐しています。

• 2番目の最小値の求め方：最小値の位置 \(i_1\) を求めた後、min(dp[:i1]) と min(dp[i1+1:]) の小さい方を取ります。\(i_1\) が端（\(0\) または \(K-1\)）のときは片側が空になるので、min() が空リストでエラーにならないよう場合分けしておきます。

  • もちろん「1 回のループで最小値と2番目の最小値を同時に求める」実装でも構いませんが、Python では min や index などの組み込み関数（C 実装）を使う方が高速です。

• 入出力の高速化：\(NK \leq 10^6\) 個の整数を読むので、Python では input() を繰り返すのではなく sys.stdin.buffer.read().split() で一括読み込みし、必要な部分だけ int に変換していきます（このコードでは各行を処理するタイミングで map(int, ...) を適用し、無駄な変換や大きな中間リストを避けています）。

• オーバーフロー：コストは最大 \(10^9 \times 10^5 = 10^{14}\) 程度になるため、C++ などでは long long を使う必要があります（Python では多倍長整数なので不要）。

• リスト内包表記 [v + m1 for v in vals] のように一括で加算し、最小値の位置だけ後から上書きすると簡潔かつ高速です。

  ソースコード

import sys

def main():
    data = sys.stdin.buffer.read().split()
    n = int(data[0]); k = int(data[1])
    pos = 2
    dp = list(map(int, data[pos:pos + k]))
    pos += k

    if n == 1:
        sys.stdout.write(str(min(dp)) + "\n")
        return

    for _ in range(n - 1):
        m1 = min(dp)
        i1 = dp.index(m1)
        if i1 == 0:
            m2 = min(dp[1:])
        elif i1 == k - 1:
            m2 = min(dp[:k - 1])
        else:
            a = min(dp[:i1])
            b = min(dp[i1 + 1:])
            m2 = a if a < b else b

        vals = list(map(int, data[pos:pos + k]))
        pos += k
        ndp = [v + m1 for v in vals]
        ndp[i1] = vals[i1] + m2
        dp = ndp

    sys.stdout.write(str(min(dp)) + "\n")

main()

---

この解説は claude5opus-high によって生成されました。