問題
数直線上に $N$ 軒の家があり、$i$ 軒目の座標は $x_i$(昇順、すべて異なる)。郵便局を、家と同じ位置にちょうど $K$ か所建てる。各家は最も近い郵便局まで歩く。歩く距離の総和の最小値を求めよ。
入力形式
N K
x_1 x_2 ... x_N
制約
| # | 条件 |
|---|---|
| 1 | $1\le K\le N\le5000$、$K\le100$ |
| 2 | $1\le x_1<x_2<\dots<x_N\le10^9$ |
| 3 | 入力はすべて整数 |
入出力例
入力例1
5 2
1 2 3 10 11出力例1
3($\{1,2,3\}$ は 2 に、$\{10,11\}$ は 10 か 11 に。$2+1=3$)
入力例2
4 1
1 2 3 10出力例2
10(1 か所なら中央値 2 か 3 に建てる。$|1-2|+0+1+8=10$)
入力例3
3 3
1 5 9出力例3
0背景と出典
「$N$ 軒を $K$ 個の連続した組に分け、各組を 1 つの郵便局(その組の中央値)で担当する」問題は IOI 2000 の Post Office(私の記憶。未確認)。$N\le300$ の元問題は $O(KN^2)$ で通る。分割統治 DP 最適化は、競技プログラミングで広く使われる名前で、条件は「最適な分割点 $\mathrm{opt}(i)$ が $i$ について単調非減少」。コストが Monge(四角不等式)を満たせば成り立つ、と私は理解している(証明・出典は未確認)。
この問題の中核は、(1) 組を連続区間と見なせること(中央値に最も近い家が担当する性質)、(2) 区間コストを $O(1)$ で出す累積和、(3) コストの Monge 性を確かめてから分割統治に載せること。(3) は下の検証ログで、乱数列 20 本での四角不等式の全数検査(合計 409500 組、違反 0 件)で確かめた。
概念図
ヒント(クリックで段階的に開示)
ヒント1(方向性)
家を左から数えて $i$ 軒目までを $k$ 個の組に分けたときの最小値を $dp_k[i]$ とする。最後の組が $(j,i]$ の家なら、そのコストは $dp_{k-1}[j]+\mathrm{cost}(j,i)$。$\mathrm{cost}(l,r)$ は区間 $[l,r)$ の家を中央値 1 点で担当するときの距離の総和で、累積和から $O(1)$ で出せる。素朴には $O(KN^2)$。
ヒント2(アプローチ)
$\mathrm{opt}_k(i)$($dp_k[i]$ を与える最小の $j$)は $i$ について単調に増える。したがって、$i$ の真ん中 $\mathrm{mid}$ の最適な $j$ を全探索で 1 つ見つければ、左半分は $j$ 以下だけ、右半分は $j$ 以上だけを探せばよい。
ヒント3(誘導)
solve(lo, hi, optlo, opthi): $\mathrm{mid}=(lo+hi)/2$ について $j\in[\mathrm{optlo},\min(\mathrm{mid}-1,\mathrm{opthi})]$ を全部試して最小の $j^*$ を取り、solve(lo, mid-1, optlo, j*) と solve(mid+1, hi, j*, opthi) を呼ぶ。各段で試す $j$ の総数が $O(N)$、段数が $O(\log N)$。
模範解答 (Python)
模範解答を表示する(まず自力で解いてから)
import sys
def main():
data = sys.stdin.buffer.read().split()
n, k = int(data[0]), int(data[1])
x = [int(v) for v in data[2:2 + n]] # 昇順
s = [0] * (n + 1)
for i in range(n):
s[i + 1] = s[i] + x[i]
def cost(l, r):
# 区間 [l, r) を 1 つの郵便局(中央値)で担当する距離の総和
m = (l + r) >> 1
return x[m] * (m - l) - (s[m] - s[l]) + (s[r] - s[m]) - x[m] * (r - m)
INF = float("inf")
prev = [INF] * (n + 1)
prev[0] = 0
for _ in range(k):
cur = [INF] * (n + 1)
def solve(lo, hi, optlo, opthi):
# cur[lo..hi] を、最適な分割点が optlo..opthi にあることを使って求める
if lo > hi:
return
mid = (lo + hi) >> 1
best = INF
bj = optlo
for j in range(optlo, min(mid - 1, opthi) + 1):
v = prev[j] + cost(j, mid)
if v < best:
best = v
bj = j
cur[mid] = best
solve(lo, mid - 1, optlo, bj)
solve(mid + 1, hi, bj, opthi)
solve(1, n, 0, n - 1)
prev = cur
print(prev[n])
main()
計算量: $O(KN\log N)$。再帰の深さは $\log_2N\approx13$ なので再帰上限の心配は無い。実測(Python 3.14.7、M1 Mac)で $N=5000,K=50$ が 0.86 s、$N=5000,K=100$ が 1.61 s、$N=3000,K=50$ が 0.5 s。
Step-by-Step 解説
中央値 $x_m$($m=\lfloor(l+r)/2\rfloor$)に局を置くと、左側の家は $x_m-x_i$、右側は $x_i-x_m$ の距離。累積和 $s$ で $\sum$ を出す: $x_m(m-l)-(s_m-s_l)+(s_r-s_m)-x_m(r-m)$。偶数個のとき中央値は 2 つあり、どちらに置いても同じ総和になる。
各家が最も近い局に行くので、局を座標順に並べると、局 $a$ の担当は連続した家の列になる(隣り合う局の中点で切る)。局の位置を家の位置に限るのは、区間内の距離の総和は中央値の家で最小になり、そこより外に置くメリットが無いため。この主張の厳密な証明は本ファイルに無いが、全分割の列挙(距離を絶対値で直接計算)と 200 件一致した(実測)。
$\mathrm{opt}_k(i)\le\mathrm{opt}_k(i+1)$。これは $\mathrm{cost}$ が Monge、$\mathrm{cost}(a,c)+\mathrm{cost}(b,d)\le\mathrm{cost}(a,d)+\mathrm{cost}(b,c)\ (a\le b<c\le d)$ のとき成り立つ(私の理解。証明は未確認)。本問のコストは、$n=25$ のランダムな昇順列 20 本で $a\le b<c\le d$ のすべての組(合計 409500 組)を調べ、違反 0 件だった(実測)。分割統治は条件が壊れても黙って誤答するので、満たす根拠を持ってから使う。
solve の 1 段(再帰の深さ 1 つ分)では、区間 $[lo,hi]$ が分かれ、探索範囲 $[\mathrm{optlo},\mathrm{opthi}]$ も分かれる(端点 $j^*$ だけを共有)。したがって 1 段で試す $j$ の総数は $O(N)$、段数は $\log N$。これが $K$ 層で $O(KN\log N)$。
右半分の下限を $j^*$ でなく $j^*+1$ にすると 300 件中 257 件で誤答(実測)。左右が端点 $j^*$ を共有するのは、$\mathrm{opt}(\mathrm{mid}+1)$ が $j^*$ と等しいことがあり得るから。平均値(切り捨て)で局を置く「近道」は、300 件中 239 件で誤答(実測)。中央値でなければ距離の総和が最小にならない。
迷った他案と採らなかった理由
| 案 | 計算量 | 採否 | 理由 |
|---|---|---|---|
| $dp_k[i]$ を全 $j$ で試す | $O(KN^2)$ | 検算用のみ | $N=5000,K=100$ で $2.5\times10^9$。Python では不可。全分割列挙との 200 件一致の検算に使った |
| 分割統治 DP 最適化(本解) | $O(KN\log N)$ | 採用 | $O(KN^2)$ と 300 件一致 |
| Knuth 最適化 | $O(N^2)$ | 不採用 | $K$ が小さい($\le100$)ので $O(KN\log N)$ の方が小さい。$K\approx N$ なら Knuth |
| Aliens trick(ラグランジュ緩和) | $O(N\log N\log\text{値域})$ | 不採用 | $K$ 層の関数 $f(K)$ が凸という前提が要る。本問は Monge から凸が導かれるはずだが、本セッションでは確認していない。$K\le100$ なら不要 |
| SMAWK で行最小を $O(N)$ | $O(KN)$ | 不採用 | 実装が重く、Python では分割統治と差が出ない見込み(未測定) |
前提が変わったら:
- **$K$ も $5000$ まで大きい**なら $O(KN\log N)=3\times10^8$ で Python では不可。Aliens trick か、Knuth(コストが Monge なので適用可)
- 郵便局の位置が家に限られない(実数)でも、距離総和は中央値で最小になるので答えは同じ。2 乗誤差(平均)なら別のコストで、Monge 性は別途検査する
- 家の座標が昇順でない入力なら、最初にソートしてから累積和を取る(本問は昇順を保証)
よくあるミス
| ミス | 原因 | 正しい書き方 |
|---|---|---|
| 右半分の下限を $j^*+1$ にする | 「$j^*$ は使い済み」と思う | solve(mid+1, hi, bj, opthi)(257/300 誤答) |
| 平均で局を置く | 最小二乗と混同する | 中央値(239/300 誤答) |
| 分割統治の条件を確かめない | 「DP の最小化なら使える」と誤解 | 四角不等式を乱択で検査(409500 組で違反 0) |
$j$ の範囲の上限を mid にする | 空の組($j=\mathrm{mid}$)を許す | min(mid - 1, opthi)(組は少なくとも 1 軒) |
次のステップ
- 発展1: Aliens trick で $K\approx N/2$ の場合を解く(凸性の確認が要る)
- 発展2: 局の位置が連続値で、距離の 2 乗の総和を最小化する版(組のコスト $\sum(x_i-\bar x)^2$)に分割統治を適用し、Monge 性を乱択で検査する
- 発展3: 分割統治 DP 最適化の正しさの証明(opt の単調性)を原典で確認する(本セッションでは未確認)
検証ログ(本セッションで実行した数値)
- 全分割列挙との照合(実測): $N\le9$ の 200 件で、全分割を列挙し距離を絶対値で直接計算した最小値と、$O(KN^2)$ の DP、プログラム全体の出力が全件一致
- $O(KN^2)$ との照合(実測): $N\le40$、座標は $[1,10^9)$ の乱数 300 件でプログラム全体の出力が全件一致
- コストの Monge 性(実測): $n=25$ の乱数昇順列 20 本、$a\le b<c\le d$ の全組(合計 409500 組)で $\mathrm{cost}(a,c)+\mathrm{cost}(b,d)>\mathrm{cost}(a,d)+\mathrm{cost}(b,c)$ となる違反は 0 件
- 誤った方法(実測): 右半分の下限を $j^*+1$ にすると 257/300、平均で局を置くと 239/300 で誤答
- 入力例(実測): 例 1
3、例 210、例 30 - 実行時間(実測): $N=5000,K=50$ が 0.86 s、$N=5000,K=100$ が 1.61 s、$N=3000,K=50$ が 0.5 s。(同じプログラムを機械が混んでいる時に測ると $N=5000,K=50$ が 4.6 s と約 5 倍に揺れた。時間は測定時の負荷に強く依存する)
- 扱っていない範囲: opt の単調性・「連続区間で足りる」ことの厳密な証明、$K>100$、Aliens trick・SMAWK との速度比較、$N>10^4$、IOI 2000 の原典確認