概要

\(N\) 枚の板を隣り合う板が異なる色になるように \(K\) 色で塗るとき、かかる総コストの最小値を求める問題です。動的計画法（DP）を状態の持ち方を工夫して高速化することで解くことができます。

考察

素朴な動的計画法 (DP)

板を左から順に塗っていくことを考えます。 \(i\) 枚目の板の色を決める際に影響するのは「直前の板（\(i-1\) 枚目）の色」のみです。 したがって、次のような動的計画法（DP）を立てることができます。

• \(DP[i][j]\): 板 \(i\) を色 \(j\) で塗る場合の、板 \(1\) から板 \(i\) までの最小コスト

遷移式は以下のようになります。 $\(DP[i][j] = C_{i, j} + \min_{k \neq j} DP[i-1][k]\)$

しかし、この遷移を愚直に計算すると、各板 \(i\) (\(1 \le i \le N\))、各色 \(j\) (\(1 \le j \le K\)) について、他のすべての色 \(k \neq j\) の最小値を求めるために \(O(K)\) かかるため、全体の計算量は \(O(N \times K^2)\) となってしまいます。\(K\) が最大 \(10^6\) なので、これでは実行時間制限に間に合いません（TLEとなります）。

高速化のアイデア（最小値と2番目の最小値）

\(\min_{k \neq j} DP[i-1][k]\) という計算に注目します。 これは「直前の板の色が \(j\) 以外であるときのコストの最小値」を求めています。

板 \(i-1\) における全色のコストのうち、 1. 最小値（コストとそのときの色） 2. 2番目に小さいコスト

の2つだけをあらかじめ求めておくとどうなるでしょうか？ - 現在の色 \(j\) が「直前の最小値の色」と異なる場合: 直前の最小値を選ぶのが最適です。 - 現在の色 \(j\) が「直前の最小値の色」と同じ場合: 隣り合う板を同じ色にできないため、直前の2番目に小さいコストを選ぶしかありません。

このように考えると、各色 \(j\) に対する遷移を \(O(1)\) で行うことができます。 さらに、板 \(i\) についても同様に「最小値」と「2番目の最小値」だけを記録して次の板の計算に進めばよいため、全色分の配列をすべて保持し続ける必要もありません。

アルゴリズム

1. 板 \(1\) のコスト \(C_{1, j}\) の中から、「最小値（値と色）」と「2番目に小さい値」を求めます。
2. \(i = 2, 3, \dots, N\) の順に、板 \(i\) を処理します。
   • 直前の板の最小値・2番目の最小値を持っておきます。
   • 各色 \(j = 1, 2, \dots, K\) について:
     • 直前の最小値の色が \(j\) でなければ、コストは \((\text{直前の最小値}) + C_{i, j}\) です。
     • 直前の最小値の色が \(j\) であれば、コストは \((\text{直前の2番目の最小値}) + C_{i, j}\) です。
     • 計算したコストをもとに、板 \(i\) の最小値・2番目の最小値を更新します。
3. 板 \(N\) まで処理が終わった後、板 \(N\) の「最小値」が求める答えです。

計算量

• 時間計算量: \(O(N K)\)
  • 各板・各色について \(O(1)\) で遷移を計算するため、全体のループ回数は \(N \times K\) 回となります。
  • 制約より \(N \times K \le 10^6\) であるため、実行時間制限に十分余裕を持って間に合います。
• 空間計算量: \(O(N K)\)
  • DPテーブル全体を保持せず、直前の状態（最小値・2番目の最小値）のみを保持しながら処理できます。
  • 入力全体を一度に読み込むために \(O(NK)\) のメモリを消費しますが、制約下で問題なく動作します。

実装のポイント

• 高速な入出力: 入力要素数が最大で約 \(10^6\) 個になるため、Pythonでは input() を繰り返し呼び出すと入出力がボトルネックになります。sys.stdin.read().split() を用いて一度に入力を読み込むことで大幅に高速化できます。

• 初期化の値: 最小値を更新していくための初期値として、十分大きな値（例えば \(10^{18}\)）を使用します。コストの総和は最大で \(10^5 \times 10^9 = 10^{14}\) になるため、\(32\) ビット整数の範囲を超える点に注意してください。

• 最小値の更新: 新しい値 cost が入ってきたとき、

  1. 最小値より小さければ、現在の最小値を2番目の最小値に降格させた上で、最小値を更新する

  2. 最小値以上かつ2番目の最小値より小さければ、2番目の最小値のみを更新する という2段階の判定を正確に行う必要があります。

     ソースコード

import sys


def solve():
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    N = int(input_data[0])
    K = int(input_data[1])

    INF = 10**18
    idx = 2

    min1_val = INF
    min1_idx = -1
    min2_val = INF

    for j in range(K):
        c = int(input_data[idx])
        idx += 1
        if c < min1_val:
            min2_val = min1_val
            min1_val = c
            min1_idx = j
        elif c < min2_val:
            min2_val = c

    for _ in range(1, N):
        next_min1_val = INF
        next_min1_idx = -1
        next_min2_val = INF

        prev_min1_val = min1_val
        prev_min1_idx = min1_idx
        prev_min2_val = min2_val

        for j in range(K):
            c = int(input_data[idx])
            idx += 1
            cost = (
                prev_min2_val + c if j == prev_min1_idx else prev_min1_val + c
            )

            if cost < next_min1_val:
                next_min2_val = next_min1_val
                next_min1_val = cost
                next_min1_idx = j
            elif cost < next_min2_val:
                next_min2_val = cost

        min1_val = next_min1_val
        min1_idx = next_min1_idx
        min2_val = next_min2_val

    print(min1_val)


if __name__ == "__main__":
    solve()

---

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