Day 081-Q2 — Knuth-Yao 高速化（四辺形不等式による $K$ 分割 DP の $O(KN \log N)$ 化）

問題

$N$ 個の石があり，それぞれ重さ $w_1, w_2, \ldots, w_N$ を持つ。隣接する石を合体させるコストは「合体後の石の重さ」。これを最終的に $K$ 個の石 が残るまで合体させる最小総コストを求めよ。

区間 $[l, r]$ を 1 つにするコスト $C[l][r]$ を以下で定義する（0-indexed）：

$N$ 個の石を $K$ 個の連続区間に分割し，各区間の $C$ の和を最小化せよ。

入出力例

5 2
1 2 3 4 5

15

6 3
3 1 4 1 5 9

18

概念図: 四辺形不等式と最適分割点の単調性

ヒント

# 分割 DP を分割統治で O(N log N) に最適化
def rec(lo, hi, opt_lo, opt_hi):
    if lo > hi: return
    mid = (lo + hi) // 2
    best = INF; best_m = opt_lo
    for m in range(opt_lo, min(opt_hi, mid)+1):
        if dp_prev[m] != INF:
            val = dp_prev[m] + C[m][mid-1]
            if val < best:
                best = val; best_m = m
    dp_cur[mid] = best
    rec(lo, mid-1, opt_lo, best_m)
    rec(mid+1, hi, best_m, opt_hi)

模範解答

import sys
input = sys.stdin.readline

def solve():
    N, K = map(int, input().split())
    w = list(map(int, input().split()))

    S = [0] * (N+1)
    for i in range(N): S[i+1] = S[i] + w[i]

    INF = float('inf')

    # C[i][j]: 区間 [i,j] を 1 つにする最小コスト (Knuth 最適化)
    C = [[0]*N for _ in range(N)]
    opt_C = [[0]*N for _ in range(N)]
    for i in range(N): opt_C[i][i] = i

    for length in range(2, N+1):
        for i in range(N - length + 1):
            j = i + length - 1
            best = INF
            lo = opt_C[i][j-1]
            hi = opt_C[i+1][j] if i+1 <= j else j
            for m in range(lo, hi+1):
                val = C[i][m] + (C[m+1][j] if m+1 <= j else 0)
                if val < best:
                    best = val; opt_C[i][j] = m
            C[i][j] = best + (S[j+1] - S[i])

    # 分割 DP
    dp_prev = [INF] * (N+1)
    dp_prev[0] = 0
    for i in range(1, N+1):
        dp_prev[i] = C[0][i-1]

    for k in range(2, K+1):
        dp_cur = [INF] * (N+1)

        def rec(lo, hi, opt_lo, opt_hi):
            if lo > hi: return
            mid = (lo + hi) // 2
            best = INF; best_m = opt_lo
            for m in range(opt_lo, min(opt_hi, mid)+1):
                if m >= k-1 and dp_prev[m] != INF:
                    val = dp_prev[m] + C[m][mid-1]
                    if val < best:
                        best = val; best_m = m
            dp_cur[mid] = best
            rec(lo, mid-1, opt_lo, best_m)
            rec(mid+1, hi, best_m, opt_hi)

        rec(k, N, k-1, N)
        dp_prev = dp_cur

    print(dp_prev[N])

solve()

Step-by-Step 解説

Step 1: 問題の DP 定式化

$C[i][j]$ = 区間 $[i, j]$ を 1 つにする最小コスト：$C[i][j] = \min_{m}(C[i][m] + C[m+1][j]) + S[j+1] - S[i]$

$dp[k][i]$ = $A[0, \ldots, i-1]$ を $k$ 個に分ける最小コスト：$dp[k][i] = \min_{m}(dp[k-1][m] + C[m][i-1])$

Step 2: 四辺形不等式の確認

$C[a][c] + C[b][d] \le C[a][d] + C[b][c]$ ($a \le b \le c \le d$) が帰納法で示せる → 最適分割点の単調性が成立。

Step 3: Knuth 最適化 ($C$ の計算を $O(N^2)$ に）

$opt\_C[i][j-1] \le opt\_C[i][j] \le opt\_C[i+1][j]$ を利用して探索範囲を絞る。全体の比較回数 $O(N^2)$。

Step 4: 分割 DP の分割統治最適化 ($O(N \log N)$ per layer)

rec(lo, hi, opt_lo, opt_hi) で二分して処理。各深さで合計 $O(N)$ 比較 × $O(\log N)$ 深さ = $O(N \log N)$ / layer。

よくあるミス

次のステップ

• 発展問題: LARSSON 型ナップサック分割 DP（WQS 二分探索 + Knuth-Yao 二重最適化）
• ACL の Convex Hull Trick との組み合わせ（傾き単調 + 分割点単調の同時適用）

自己評価