問題
昇順に並んだ $N$ 個の鍵があり、$i$ 番目の鍵が検索される回数は $p_i$ である。これらを二分探索木にする(左の子孫は小さい鍵、右の子孫は大きい鍵)。鍵 $i$ の深さを、根を 1 として $d_i$ とするとき、
$$\text{コスト}=\sum_{i=1}^{N} p_i\,d_i$$
を最小にする木のコストを求めよ。
入力形式
N
p_1 p_2 ... p_N
制約
| # | 条件 |
|---|---|
| 1 | $1\le N\le2000$ |
| 2 | $1\le p_i\le10^9$ |
| 3 | 入力はすべて整数 |
入出力例
入力例1
3
10 20 30出力例1
100(根を鍵 2 にすると $20\cdot1+10\cdot2+30\cdot2=100$。根を鍵 3 の鎖にすると $30+40+30=100$ のように同点の木もある)
入力例2
1
5出力例2
5入力例3
5
4 2 6 3 5出力例3
39背景と出典
最適二分探索木の $O(N^2)$ 解は、私の記憶では Knuth(1971)の "Optimal binary search trees" による。「最適な根は $\mathrm{opt}[l][r-1]\le\mathrm{opt}[l][r]\le\mathrm{opt}[l+1][r]$ を満たす」ことを使う。この単調性を、コスト関数の四角不等式(Monge 性)から一般に導いたのは Yao(1980、同じく記憶)。本セッションでは両方とも文献を開いておらず、年・題名・帰属は未確認。
正しさの根拠は証明ではなく、$O(N^3)$ の素朴な DP との 400 件の一致と、条件を破ったときの誤り率の実測。「$\min_k$ の範囲が狭まる」ことを使う高速化は、条件(四角不等式)が成り立つときにしか正しくない。成り立たないコストに同じコードを当てると 300 件中 259 件で誤答した(実測)。エラーにならず、それらしい値を返すのが怖い所。
概念図
ヒント(クリックで段階的に開示)
ヒント1(方向性)
区間 $[l,r)$ の鍵だけで作る木のコスト $dp[l][r]$ を考える。根を鍵 $k$ にすると、左右の部分木は $[l,k)$ と $[k+1,r)$。全体の深さが 1 つずつ増えるので、コストは $dp[l][k]+dp[k+1][r]+\sum_{i\in[l,r)}p_i$。$k$ を全部試すと $O(N^3)$。
ヒント2(アプローチ)
最適な根 $\mathrm{opt}[l][r]$ を覚えておく。区間の右端を 1 つ伸ばしても($[l,r-1)\to[l,r)$)、左端を 1 つ縮めても($[l,r)\to[l+1,r)$)、最適な根は逆戻りしない。つまり $\mathrm{opt}[l][r-1]\le\mathrm{opt}[l][r]\le\mathrm{opt}[l+1][r]$。
ヒント3(誘導)
区間の長さが短いものから順に埋める。$dp[l][r]$ を求めるとき、$k$ を $\mathrm{opt}[l][r-1]$ から $\mathrm{opt}[l+1][r]$ までだけ試す。同じ長さの区間をずらして足し合わせると、試す $k$ の総数が望遠鏡和になって $O(N)$ ずつ、全長で $O(N^2)$。
模範解答 (Python)
模範解答を表示する(まず自力で解いてから)
import sys
from array import array
def main():
data = sys.stdin.buffer.read().split()
n = int(data[0])
p = [int(x) for x in data[1:1 + n]]
s = [0] * (n + 1)
for i in range(n):
s[i + 1] = s[i] + p[i]
# dp[l][r]: 区間 [l, r) の鍵だけで作る最小コスト。opt[l][r]: その最適な根
dp = [array('q', [0]) * (n + 1) for _ in range(n + 1)]
opt = [array('i', [0]) * (n + 1) for _ in range(n + 1)]
for l in range(n):
opt[l][l + 1] = l
dp[l][l + 1] = p[l]
for length in range(2, n + 1):
for l in range(0, n - length + 1):
r = l + length
lo = opt[l][r - 1] # Knuth: opt[l][r-1] <= opt[l][r] <= opt[l+1][r]
hi = opt[l + 1][r]
dl = dp[l]
best = dl[lo] + dp[lo + 1][r]
bk = lo
for k in range(lo + 1, hi + 1):
v = dl[k] + dp[k + 1][r]
if v < best:
best = v
bk = k
dp[l][r] = best + s[r] - s[l]
opt[l][r] = bk
print(dp[0][n])
main()
計算量: $O(N^2)$ 時間・$O(N^2)$ メモリ。実測(Python 3.14.7、M1 Mac)で $N=1000$ が 0.3 s、$N=2000$ が 1.0 s・最大メモリ 66 MB(array で持った場合)。
Step-by-Step 解説
$dp[l][r]=\min_{l\le k<r}\bigl(dp[l][k]+dp[k+1][r]\bigr)+W(l,r)$、$W(l,r)=\sum_{i\in[l,r)}p_i$。空区間は 0。区間は半開 $[l,r)$ で持つと、$k$ を根にしたときの左右が $[l,k)$ と $[k+1,r)$ で境界の足し引きが要らない。素朴な実装は $O(N^3)$($N=2000$ で $1.3\times10^9$ 回)。
$\mathrm{opt}[l][r]$ を最小を与える最小の $k$ とする。右端を伸ばした区間 $[l,r)$ の最適根は $[l,r-1)$ のそれより右、左端を縮めた $[l+1,r)$ のそれより左。この 2 つで挟まれた範囲だけを試す。
同じ長さ $L$ の区間 $[l,l+L)$ を $l=0,1,\dots$ の順に見ると、$l$ の試す範囲は $[\mathrm{opt}[l][r-1],\ \mathrm{opt}[l+1][r]]$ で、次の $l+1$ の範囲は $\mathrm{opt}[l+1][r]$ から始まる。隣り合う範囲が端点を共有するので、長さ $L$ の区間全体で試す $k$ の総数は $O(N+L)=O(N)$。長さ $L$ は $N$ 通りなので全体 $O(N^2)$。実測(内側のループの総回数)で $N=200$ が 39529 回、$N=400$ が 160040 回。$N^2$ は 40000 と 160000 でほぼ一致、素朴な範囲 $N^3/6$ は 1333333 と 10666666(実測)。
この単調性は $W$ が四角不等式 $W(a,c)+W(b,d)\le W(a,d)+W(b,c)\ (a\le b\le c\le d)$ と単調性 $W(b,c)\le W(a,d)$ を満たすことから導かれる(私の記憶する Yao の結果。証明は本ファイルに無い)。区間の和 $W(l,r)=\sum p_i$ は等号で満たす。満たさないコストに同じコードを当てると、300 件中 259 件で誤答した(実測)。一方 $W=(\sum p_i)^2$(四角不等式を満たす)では 300 件すべて素朴 DP と一致した(実測)。
下限を $\mathrm{opt}[l][r-1]+1$ にすると 300 件中 280 件、左の最適根だけを使う(範囲を 1 点にする)と 279 件で誤答(実測)。範囲の端を含むかどうかで答えが変わるので、$\le$ の両側を含めて試す。
迷った他案と採らなかった理由
| 案 | 計算量 | 採否 | 理由 |
|---|---|---|---|
| 全 $k$ を試す区間 DP | $O(N^3)$ | 検算用のみ | $N=2000$ で $1.3\times10^9$ 回。Python では不可 |
| Knuth 最適化(本解) | $O(N^2)$ | 採用 | 素朴 DP と 400 件一致。条件が成り立つことを問題側で保証できる |
| Garsia–Wachs / Hu–Tucker | $O(N\log N)$ | 不採用 | これは葉だけに重みがある木(アルファベット木)の最適化。本問は内部節点にも重みがある別の問題で、使えない(Day 158 Q3 は Garsia–Wachs) |
| 分割統治 DP 最適化 | $O(N^2\log N)$ | 不採用 | 区間 DP では opt が 2 次元に動くので、Knuth の方が 1 段速い |
array でなく list of list | 同じ | 不採用 | $N=2000$ で 400 万要素の Python int は 100 MB 超(概算)。array('q') で 66 MB(実測) |
前提が変わったら:
- **検索に失敗する確率 $q_i$(鍵と鍵の間)も入れる**(教科書版)なら、$W(l,r)$ に $q$ の和を足すだけで四角不等式は保たれる(概算。未検証)
- コストが区間の和でない(たとえば区間の分散)なら、四角不等式を乱数で確かめてから使う。確かめずに使うのが最も危険
- **$N\ge10^4$**なら $O(N^2)$ でも遅い。$p_i$ が単調などの追加構造を探す
よくあるミス
| ミス | 原因 | 正しい書き方 |
|---|---|---|
| 条件を確かめずに Knuth を使う | 「区間 DP は Knuth」と覚える | 四角不等式を満たすかを小さい例で確かめる(非 QI コストで 259/300 誤答) |
| 下限を $\mathrm{opt}[l][r-1]+1$ にする | 同点のとき 2 回目を省きたくなる | 両端を含める(280/300 誤答) |
opt[l][r-1] だけを根にする | 「根は動かない」と誤解 | 範囲 $[\mathrm{opt}[l][r-1],\ \mathrm{opt}[l+1][r]]$ を試す(279/300 誤答) |
| 長さの小さい順でなく $l$ の降順で全 $r$ を埋める | opt[l+1][r] が未計算になる | 長さの外側ループ、または $l$ 降順・$r$ 昇順 |
次のステップ
- 発展1: 失敗確率 $q_i$ を含む教科書版(Cormen らの形)に拡張し、素朴 DP と比べる
- 発展2: 四角不等式を乱数テストで判定する関数を書き、本問の $W$ と、わざと壊した $W$ で違反数を数える
- 発展3: Knuth–Yao の定理(Monge 性から単調性が出る)の証明を、原典で確認する(本セッションでは未確認)
検証ログ(本セッションで実行した数値)
- 素朴 DP との照合(実測): $N\le12$、$p_i\in[1,3]$・$[1,20]$・$[1,10^9]$ の混合 400 件で、プログラム全体の出力が $O(N^3)$ の再帰 DP と全件一致
- 誤った方法(実測): 下限 $\mathrm{opt}[l][r-1]+1$ は 280/300、左の根だけ使うのは 279/300 で誤答。四角不等式を満たさないランダムなコスト $W$($0$〜$50$ の乱数)に Knuth を当てると 259/300 誤答。$W=(\sum p_i)^2$ では 300 件中 0 件の誤答
- 内側ループの総回数(実測): $N=200$ で 39529、$N=400$ で 160040($N^2$ は 40000、160000)
- 入力例(実測): 例 1
100、例 25、例 339 - 実行時間・メモリ(実測): $N=1000$ が 0.3 s、$N=1500$ が 0.6 s、$N=2000$ が 1.0 s・66 MB。$N=3000$ は 2.4 s・135 MB
- 扱っていない範囲: 四角不等式から単調性が出る証明、失敗確率 $q_i$ 付きの版、$N>3000$、Knuth・Yao の原典確認、木そのものの復元