Official

D - フェンスの塗装 / Painting the Fence Editorial by admin Editorial for printing

gpt-6-astra-medium

概要

隣り合う板の色が異なるように塗るときの、総コストの最小値を求めます。動的計画法(DP)を使い、直前の板までの「最小値」と「2 番目の最小値」を管理することで、\(O(NK)\) 時間で解けます。

考察

最後に塗った色を状態にする

板 \(i\) に使える色は、板 \(i-1\) の色によって決まります。そこで、次の DP を考えます。

\(dp[i][j] =\) 板 \(1\) から板 \(i\) までを条件を満たすように塗り、板 \(i\) の色を \(j\) にしたときの総コストの最小値

最初の板には隣の板との制約がないため、\(dp[1][j] = C_{1,j}\) です。

板 \(i\) を色 \(j\) で塗る場合、直前の板には \(j\) 以外の色を使う必要があるので、\(i \geq 2\) では次のように遷移できます。

\(dp[i][j] = C_{i,j} + \min_{x \neq j} dp[i-1][x]\)

ただし、各状態で \(j\) 以外の全色を調べると、時間計算量は \(O(NK^2)\) となり、制限時間に間に合いません。

最小値と 2 番目の最小値だけを覚える

遷移で必要なのは、「ある 1 色を除いたときの最小値」です。

直前の DP について、次の 3 つを管理します。

  • best:全色の中での最小値
  • best_color:その最小値を取る色の 1 つ
  • second:best_color を除いた色の中での最小値

現在の板を色 color で塗るときは、次のように判断できます。

  • color != best_color なら、best_color から遷移できるので best を使う
  • color == best_color なら、その色からは遷移できないので second を使う

これにより、1 つの状態を \(O(1)\) 時間で計算できます。

例えば、直前の DP が色 \(1,2,3\) の順に \((4,4,9)\) なら、best = 4、best_color = 1、second = 4 とできます。色 \(1\) を除いても、色 \(2\) の値 \(4\) が使えるためです。

ここでの「2 番目の最小値」は、最小値より大きい値のうち最小のものではありません。異なる色が同じ値を持つ場合は、best と second が等しくなることもあります。

アルゴリズム

コードでは、DP 配列そのものは持たず、上記の 3 つの情報だけを更新していきます。

  1. best = second = 0、best_color = -1 として初期化します。

    • -1 は実際には存在しない色です。
    • 最初の板ではすべての色について best = 0 が使われるので、計算される値はそのまま \(C_{1,j}\) になります。
  2. 板を左から順に処理します。

    • 現在の板についての最小値 new_best と 2 番目の最小値 new_second を、十分大きな値で初期化します。
    • 各色 color について、次の値を計算します。

    value = cost + (second if color == best_color else best)

    • value を使って、現在の板についての情報を更新します。
      • value < new_best なら、これまでの new_best を new_second に移し、new_best とその色を更新します。
      • そうでなく、value < new_second なら、new_second のみを更新します。
  3. その板の全色を処理したら、best、second、best_color を、新しく求めた情報に置き換えます。

  4. すべての板を処理した後の best を出力します。

各色の value は、直前の板を異なる色で塗った場合の最小コストに、現在の塗装コストを加えたものです。したがって、この更新は DP の遷移式を正しく計算しています。最後の best は、最後の板の色を問わない最小コストとなります。

計算量

  • 時間計算量: \(O(NK)\)
    各板の各色について、定数時間の計算を行います。制約より \(NK \leq 10^6\) なので、十分高速です。
  • 空間計算量: \(O(NK)\)
    提示コードは入力を read().split() で一括読み込みするため、その保持に \(O(NK)\) の領域を使います。DP の計算自体に必要な追加領域は \(O(1)\) です。

実装のポイント

  • 同じ値を持つ別の色も、2 番目の候補に含めます。
    コードでは value == new_best の場合でも、続く elif value < new_second によって new_second に登録できます。

  • 直前の板の情報と現在の板の情報を分けます。
    色を調べている途中で best などを書き換えると、別の色の計算に現在の板の情報が混ざってしまいます。new_best などに集計し、全色を処理してから置き換えます。

  • 大きな整数を扱います。
    総コストは最大 \(N \times 10^9 \leq 10^{14}\) になります。Python ではそのまま扱えますが、固定長整数の言語では 64 bit 整数が必要です。コードの inf = 10**30 は、答えより十分大きな値です。

  • コード内の色番号は \(0\) から \(K-1\) です。初期値の best_color = -1 は、どの実際の色とも一致しません。

    ソースコード

import sys


def main():
    data = map(int, sys.stdin.buffer.read().split())
    n = next(data)
    k = next(data)

    inf = 10**30
    best = second = 0
    best_color = -1

    for _ in range(n):
        new_best = inf
        new_second = inf
        new_color = -1

        for color in range(k):
            cost = next(data)
            value = cost + (second if color == best_color else best)

            if value < new_best:
                new_second = new_best
                new_best = value
                new_color = color
            elif value < new_second:
                new_second = value

        best = new_best
        second = new_second
        best_color = new_color

    print(best)


if __name__ == "__main__":
    main()

この解説は gpt-6-astra-medium によって生成されました。

posted:
last update: