Day 082-Q4 — 区間DP + 凸包トリック融合（区間最小コスト分割 $O(N^2 \to KN \log N)$）

問題

長さ $N$ の数列 $a_1, \ldots, a_N$ を $K$ 個の連続する区間に分割する。各区間 $[l, r]$ のコストを $\text{cost}(l, r) = (S_r - S_{l-1})^2$ とする（$S_i = \sum_{j=1}^{i} a_j$）。コストの総和を最小化せよ。

入出力例

6 3
3 1 4 1 5 9

5 2
1 2 3 4 5

97

85

概念図: CHT による DP 最適化

ヒント

for j in range(1, K+1):
    lct = LCTree(0, S[N]+1)
    lct.add(-2 * S[0], prev[0] + S[0]**2)  # l=0 を追加
    for i in range(1, N+1):
        val = lct.query(S[i])               # dp[j][i] = val + S[i]^2
        curr[i] = val + S[i]**2 if val < inf else inf
        if prev[i] < inf:
            lct.add(-2 * S[i], prev[i] + S[i]**2)  # l=i を追加
    prev = curr

模範解答

import sys
from math import inf
input = sys.stdin.readline

def solve():
    N, K = map(int, input().split())
    a = list(map(int, input().split()))
    S = [0] * (N+1)
    for i in range(N): S[i+1] = S[i] + a[i]

    MAX_VAL = S[N] + 1

    class LCTree:
        def __init__(self, lo, hi):
            self.lo = lo; self.hi = hi
            self.lines = {}

        def _add(self, node, lo, hi, m, b):
            mid = (lo + hi) >> 1
            if node not in self.lines:
                self.lines[node] = (m, b); return
            om, ob = self.lines[node]
            left_new = m * lo + b < om * lo + ob
            mid_new  = m * mid + b < om * mid + ob
            if mid_new:
                self.lines[node] = (m, b); m, b = om, ob
            if lo == hi: return
            if left_new != mid_new:
                self._add(2*node, lo, mid, m, b)
            else:
                self._add(2*node+1, mid+1, hi, m, b)

        def add(self, m, b):
            self._add(1, self.lo, self.hi, m, b)

        def _query(self, node, lo, hi, x):
            if node not in self.lines: return inf
            res = self.lines[node][0]*x + self.lines[node][1]
            if lo == hi: return res
            mid = (lo + hi) >> 1
            if x <= mid:
                return min(res, self._query(2*node, lo, mid, x))
            else:
                return min(res, self._query(2*node+1, mid+1, hi, x))

        def query(self, x):
            return self._query(1, self.lo, self.hi, x)

    prev = [inf] * (N+1)
    prev[0] = 0

    for j in range(1, K+1):
        lct = LCTree(0, MAX_VAL)
        lct.add(-2 * S[0], prev[0] + S[0]**2)
        curr = [inf] * (N+1)
        for i in range(1, N+1):
            val = lct.query(S[i])
            if val < inf:
                curr[i] = val + S[i]**2
            if prev[i] < inf:
                lct.add(-2 * S[i], prev[i] + S[i]**2)
        prev = curr

    print(prev[N])

solve()

Step-by-Step 解説

Step 1: DP の定式化

$dp[j][i]$: 最初の $i$ 要素を $j$ 区間に分割した最小コスト。

Step 2: 遷移の変形（CHT 適用準備）

$(S_i - S_l)^2 = S_i^2 - 2 S_i S_l + S_l^2$ と展開すると：

Step 3: Li Chao Tree での高速化

$l$ を増やしながら直線 $(m_l, b_l)$ を Li Chao Tree に追加し、各 $i$ でクエリを $O(\log C)$ で処理。全体 $O(KN \log C)$。$S_i$ が単調増加なら傾き単調版 CHT（deque）で $O(KN)$ も可能。

Step 4: 空間最適化

prev = [inf] * (N+1); prev[0] = 0
for j in range(1, K+1):
    curr = [inf] * (N+1)
    # ... 計算後 prev = curr
# prev[N] が答え

よくあるミス

次のステップ

• 発展問題: Aliens Trick（WQS 二分探索）で $K$ をパラメータ化し $O(N \log N \log C)$
• 参考: $\text{cost}(l,r) = c \cdot (r-l+1)$ との複合コスト（Knuth 最適化との組み合わせ）

自己評価