問題
$N$個の整数列 $a_1,\dots,a_N$ を、ちょうど$K$個の空でない連続区間に分割する。各グループのコストを「含まれる要素の総和の2乗」とし、全グループのコスト総和を最小化せよ。
$\mathrm{prefix}(i)=a_1+\dots+a_i$として、区間$(j,i]$のコストは$\mathrm{cost}(j,i)=(\mathrm{prefix}(i)-\mathrm{prefix}(j))^2$。この関数は四角不等式を満たすため、$dp[k][i]=\min_{j入力形式
N K
a_1 a_2 ... a_N制約
$1 \le K \le N \le 2\times10^5$
$0 \le a_i \le 10^4$
入出力例
入力例1
5 2
1 2 3 4 5出力例1
117{1,2,3}と{4,5}に分割：6²+9²=36+81=117 が最小。
概念図：探索範囲の絞り込み
ヒント（段階的開示）
ヒント1（方向性）
素朴な$dp[k][i]=\min_{j
ヒント2（アプローチ）
区間$[lo,hi]$の中央$mid$の最適$j$を$[optlo,opthi]$の範囲だけ探索して求め、その結果を軸に左半分の探索範囲は$[optlo,\mathrm{opt}(mid)]$、右半分は$[\mathrm{opt}(mid),opthi]$に絞る。各$j$が調べられる回数はならし$O(\log N)$回に抑えられる。
ヒント3（誘導）
def solve_layer(dp_prev, N, k, cost):
dp_cur = [float('inf')] * (N + 1)
def rec(lo, hi, opt_lo, opt_hi):
if lo > hi:
return
mid = (lo + hi) // 2
best_val, best_j = float('inf'), opt_lo
for j in range(opt_lo, min(mid - 1, opt_hi) + 1):
if dp_prev[j] == float('inf'):
continue
v = dp_prev[j] + cost(j, mid)
if v < best_val:
best_val, best_j = v, j
dp_cur[mid] = best_val
rec(lo, mid - 1, opt_lo, best_j)
rec(mid + 1, hi, best_j, opt_hi)
rec(k, N, k - 1, N - 1)
return dp_cur模範解答 (Python)
import sys
def main():
data = sys.stdin.buffer.read().split()
idx = 0
N = int(data[idx]); idx += 1
K = int(data[idx]); idx += 1
a = [int(x) for x in data[idx:idx + N]]; idx += N
prefix = [0] * (N + 1)
for i in range(N):
prefix[i + 1] = prefix[i] + a[i]
def cost(j, i):
s = prefix[i] - prefix[j]
return s * s
INF = float('inf')
dp_prev = [INF] * (N + 1)
dp_prev[0] = 0
for k in range(1, K + 1):
dp_cur = [INF] * (N + 1)
def rec(lo, hi, opt_lo, opt_hi):
if lo > hi:
return
mid = (lo + hi) // 2
best_val, best_j = INF, opt_lo
hi_j = min(mid - 1, opt_hi)
for j in range(opt_lo, hi_j + 1):
if dp_prev[j] == INF:
continue
v = dp_prev[j] + cost(j, mid)
if v < best_val:
best_val, best_j = v, j
dp_cur[mid] = best_val
rec(lo, mid - 1, opt_lo, best_j)
rec(mid + 1, hi, best_j, opt_hi)
rec(k, N, k - 1, N - 1)
dp_prev = dp_cur
print(dp_prev[N])
main()Step-by-Step 解説
1DPの定式化
$dp[k][i]$=最初の$i$要素をちょうど$k$グループに分割したときの最小コスト総和。素朴には各層$O(N^2)$。
$dp[k][i]$=最初の$i$要素をちょうど$k$グループに分割したときの最小コスト総和。素朴には各層$O(N^2)$。
2四角不等式と単調性
$\mathrm{cost}(j,i)$は四角不等式を満たすため最適$j=\mathrm{opt}(i)$は$i$の増加に対して単調非減少になる。
$\mathrm{cost}(j,i)$は四角不等式を満たすため最適$j=\mathrm{opt}(i)$は$i$の増加に対して単調非減少になる。
3分割統治による絞り込み
中央の最適$j$を軸に左右の探索範囲を分割し、各層の探索総量をならし$O(N\log N)$に抑える。
中央の最適$j$を軸に左右の探索範囲を分割し、各層の探索総量をならし$O(N\log N)$に抑える。
4層をまたぐ計算量
K層それぞれ$O(N\log N)$のため全体$O(NK\log N)$。
K層それぞれ$O(N\log N)$のため全体$O(NK\log N)$。
5具体例のトレース
$N=5,K=2$で$k=1$層は$dp[1][i]=\mathrm{prefix}(i)^2$、$k=2$層で$dp[2][5]=117$（j=3、{1,2,3},{4,5}）が最小。
$N=5,K=2$で$k=1$層は$dp[1][i]=\mathrm{prefix}(i)^2$、$k=2$層で$dp[2][5]=117$（j=3、{1,2,3},{4,5}）が最小。
よくあるミス
| ミス | 原因 | 正しい書き方 |
|---|---|---|
| 探索範囲の上限をmin(mid,opt_hi)にしてしまう | j=midを許すと右側グループが空になる | min(mid-1,opt_hi)としてj<midを保証 |
| コスト関数が四角不等式を満たすか確認せず適用 | 単調性が成立しない一般関数では最適解を見落とす | 適用前にコスト関数の凸性・四角不等式を確認 |
| dp_prev[j]がINF(未到達)のケースを無視 | kグループには最低k要素必要で範囲外jは無効 | dp_prev[j]==INFのときはスキップ |
| 再帰の基準区間をkに依存しない形で渡す | kグループではi≥k, j≥k-1でなければ無意味 | rec(k, N, k-1, N-1)のように層ごとの有効範囲を渡す |
| 再帰関数をループ内で毎回defしオーバーヘッド増 | 関数呼び出しコストを軽視 | 大規模データでは非再帰化やPyPy/C++移植を検討 |
次のステップ
- 発展: コスト関数の凸性が保証されない区間DP向けのKnuthの最適化との違いを整理せよ。
- 発展2: Kが大きい場合に有効なAliens Trick (WQS binary search)で$O(N\log N\log(\text{value range}))$に落とす手法と比較せよ。