問題
$N$ 日分の株価 $p_1,\dots,p_N$ が分かっている。株は常に 0 株か 1 株しか持てない。「1 株を買い、後でその株を売る」を 1 回の取引と数え、**取引は高々 $K$ 回**まで。同じ日に売ってから買い直してもよい。利益(売値の総和 − 買値の総和)の最大値を求めよ(取引しなければ 0)。
入力形式
N K
p_1 p_2 ... p_N
制約
| # | 条件 |
|---|---|
| 1 | $1\le N\le2\times10^5$ |
| 2 | $0\le K\le N$ |
| 3 | $0\le p_i\le10^9$ |
| 4 | 入力はすべて整数 |
入出力例
入力例1
6 2
3 2 6 5 0 3出力例1
7(2 で買って 6 で売り(4)、0 で買って 3 で売る(3)。合計 7)
入力例2
7 2
1 2 4 2 5 7 2出力例2
8(1 で買って 4 で売り(3)、2 で買って 7 で売る(5)。合計 8。1 で買って 7 で売る 1 回だけだと 6)
入力例3
3 1
5 4 3出力例3
0背景と出典
「回数 $k$ の最適値 $g(k)$ が上に凸なら、回数を **ペナルティ $\lambda$ 円/回** の手数料に置き換え、$\lambda$ を二分探索して $K$ 回に合わせる」手法は、日本の競プロで Aliens trick(IOI 2016 の問題 "Aliens" に由来すると私は記憶している。本セッションでは原典・解説を開いておらず未確認)、Lagrange 緩和、WQS 二分探索などと呼ばれる。
この問題の形(株の売買を高々 $K$ 回)は LeetCode の "Best Time to Buy and Sell Stock IV" と同型と記憶している(未確認)。$N$ が小さいときは $O(NK)$ の DP が普通の解で、本問は $N,K$ が $2\times10^5$ なのでその DP が使えない。正しさの根拠は $O(NK)$ DP との照合(下の検証ログ)。**$g(k)$ が上に凸であることは証明していない**(実測で 500 件の照合が通ったことが根拠)。
概念図
ヒント(クリックで段階的に開示)
ヒント1(方向性)
$g(k)$ = 「ちょうど(高々)$k$ 回の取引での最大利益」とおくと、$k$ を増やすと増分 $g(k)-g(k-1)$ は小さくなっていく(直感: 利益の大きい取引から順に使う)。この凸性があれば、回数の制約を「1 回あたり $\lambda$ 円払う」に取り替えて、$\lambda$ を動かすと使う回数が単調に変わる。
ヒント2(アプローチ)
手数料 $\lambda$ を固定した版: 「取引ごとに $\lambda$ 円払って、回数に制限なしで最大利益」は $O(N)$ の 2 状態 DP(株を持っている/いない)。この DP が返す回数が $K$ 以下になる最小の $\lambda$ を二分探索し、答えを $\text{(手数料込みの最大利益)}+\lambda K$ で復元する。
ヒント3(誘導)
手数料込みの最大利益が同じになる回数が複数あるとき(同点)、DP が返す回数が決まらないと二分探索が壊れる。同じ利益なら回数が少ない方を選ぶとそろえる。実装は、利益を $2^{20}$ 倍して、そこから回数を 1 引いた整数で比べると 1 つの数値で済む。
模範解答 (Python)
模範解答を表示する(まず自力で解いてから)
import sys
SH = 20 # 回数を下位 20 ビットに詰める(回数 < 2^20 = 1048576)
def solve(p, lam):
"""1 回の売買ごとに lam 円の手数料。(利益の最大, その中で最小の売買回数) を 1 つの整数に詰めて返す"""
# 値 = 利益 * 2^SH - 回数。大きいほど「利益が大きく、同利益なら回数が少ない」
free = 0 # 株を持っていない
hold = -(p[0] + lam) << SH # 持っている(買った時点で手数料を払い、回数を 1 数える)
hold -= 1
for i in range(1, len(p)):
c = p[i] << SH
nf = hold + c # 売る(売っても回数は増やさない。買った時に数えた)
if nf > free:
free = nf
nh = free - (p[i] + lam << SH) - 1 # 買う
if nh > hold:
hold = nh
return free
def main():
data = sys.stdin.buffer.read().split()
n, k = int(data[0]), int(data[1])
p = [int(v) for v in data[2:2 + n]]
# k 回で足りる(上り坂の総和)場合は手数料 0 で終わり
lo, hi = 0, max(p) - min(p) + 1 # hi 円なら 1 回も売買しない
# 売買回数が k 以下になる最小の lam を探す
while lo < hi:
mid = (lo + hi) // 2
v = solve(p, mid)
cnt = -(-v >> SH) * (1 << SH) - v # v = 利益 * 2^SH - 回数 から回数を取り出す
if cnt <= k:
hi = mid
else:
lo = mid + 1
v = solve(p, lo)
cnt = -(-v >> SH) * (1 << SH) - v
val = (v + cnt) >> SH # 手数料込みの最大利益
print(val + lo * k)
main()
計算量: $O(N\log(\max p-\min p))$。実測(Python 3.14.7、M1 Mac)で $N=10^5,K=500$ が 0.51 s、$N=10^5,K=5\times10^4$ が 0.54 s、$N=2\times10^5,K=1000$ が 1.03 s。$O(NK)$ の DP は $N=20000,K=1000$ で 5.06 s、本解は同じ入力で 0.17 s。
Step-by-Step 解説
$g(k)$ が上に凸(増分が減っていく)なら、手数料 $\lambda$ のもとで「$g(k)-\lambda k$ を最大にする $k$」は $g$ の傾きが $\lambda$ を下回り始める点。$\lambda$ を大きくすると最適な $k$ は減る。傾きは整数なので、$\lambda$ は整数だけ調べればよい(株価が整数なので $g$ の値も整数)。
状態は「持っていない/持っている」。買うときに $\lambda$ を払い、売るときは何も引かない。回数は買ったときに 1 数える。これで「回数制限なしで、手数料込みの最大利益」が $O(N)$ で出る。
手数料込みの最大利益を達成する回数は、一般に区間 $[a(\lambda),b(\lambda)]$ になる。DP が最小の $a$ を返すか最大の $b$ を返すか、決まらないまま二分探索すると壊れる。本解は「同利益なら少ない回数」にそろえ(比べる値を 利益 << 20 - 回数 にする)、**$a(\lambda)\le K$ となる最小の $\lambda$** を探す。このとき $\lambda-1$ では $a(\lambda-1)>K$ だから、$K<a(\lambda-1)\le b(\lambda)$ となって $K\in[a(\lambda),b(\lambda)]$ が言える。区間内の $K$ では、手数料込みの利益と $\lambda K$ の和が $g(K)$ に等しい。
答え $=\text{(手数料込みの最大利益)}+\lambda K$。ここで **$K$ を足す**(DP が返した回数 $c$ ではない)。$c<K$ のとき $c$ を足すと小さすぎる。最大回数を返す DP のまま「$\le K$ になる最小の $\lambda$」を探した場合や、$\lambda c$ を足した場合の誤り件数は下の表。
$\lambda=0$ で回数が $K$ 以下に収まるなら($K$ が十分大きい)、そのまま利益の最大。$K=0$ なら $\lambda=\max p-\min p+1$ まで上がり、回数 0・利益 0。上の二分探索の範囲 [0, max-min+1] はどちらも含む。
迷った他案と採らなかった理由
| 案 | 計算量 | 採否 | 理由 |
|---|---|---|---|
| $O(NK)$ の DP(回数 $k$ ごとに持つ) | $O(NK)$ | 検算用のみ | $N=20000,K=1000$ で 5.06 s、$N=2\times10^5$ には不可。検算に使った |
| 上り坂(連続する上昇区間)を列挙して大きい順に $K$ 個取る | $O(N\log N)$ | 不採用 | 谷と山の間に小さな下りがある区間は、つなぐか分けるかで $K$ の使い方が変わり、貪欲では決まらない(反例は未作成、見立てのみ) |
| 優先度付きキューで「分割・併合」を戻す貪欲 | $O(N\log N)$ | 不採用 | 正しい解法として知られていると記憶しているが、本セッションで検証していない。扱う範囲を絞った |
| Aliens trick(同点は最小回数、本解) | $O(N\log V)$ | 採用 | $O(NK)$ DP と 500 件全件一致 |
| Aliens trick(同点は最大回数のまま最小 $\lambda$ を探す) | $O(N\log V)$ | 不採用 | 500 件中 41 件誤答 |
前提が変わったら:
- **ちょうど $K$ 回**を求めるなら、$g$ が単調でなくなる(取引すると損をする日があるため)。$\lambda$ の下限を負にして同じ手順を使えるが、本問では扱わない
- **$g$ が凸でない**(取引に手数料以外の制約がある、同時に複数株を持てる等)と Aliens は使えない。小さい入力で $g(k)$ を表にして増分の単調性を確かめるのが最初の作業
- **回数 $K$ が小さい**($K\le30$ 程度)なら素直な $O(NK)$ の DP が速くて安全
よくあるミス
| ミス | 原因 | 正しい書き方 |
|---|---|---|
| 同点の回数をそろえない(最大回数を返す DP で最小 $\lambda$ を探す) | 「利益が同じなら回数はどれでもよい」と思う | 同利益なら回数を最小にそろえる(41/500 誤答) |
答えを 利益 + lam * 回数 にする | $c<K$ のとき足りなくなる | 利益 + lam * K(9/500 誤答) |
| $\lambda$ を実数で二分探索する | 傾きは整数のはずなのに浮動小数で打ち切る | 整数の $\lambda$ だけ調べる |
| 回数を詰める桁数が足りない | $N=2\times10^5$ に対して SH = 10 など | $2^{\mathrm{SH}}>N$ になる SH = 20 |
次のステップ
- 発展1: 同じ手順で「ちょうど $K$ 個の互いに素な区間を選んで総和最大」を解く(回数 $=$ 区間の個数、同点の扱いは本問と同じ)
- 発展2: $g(k)$ が上に凸であることを証明する(最小費用流の増加路が 1 本ずつ流量を増やす議論に帰着できることが知られていると記憶しているが、未確認)
- 次回予告: 焦点帯外(M2・M3)の問題
検証ログ(本セッションで実行した数値)
- $O(NK)$ DP との照合(実測): $N\le12$、$K\le7$、$0\le p\le9$ の 500 件でプログラム全体が全件一致。そのうち $K\le2$ は 176 件、「取引回数に上限がなければ答えが増える」($K$ が効く)のは 114 件
- 誤った方法(実測): 最大回数を返す DP+最小 $\lambda$ 探索 41/500、
利益 + lam * 回数9/500 で誤答(最小回数+lam * Kの正しい形は 0/500) - 入力例(実測): 例 1 は 7、例 2 は 8、例 3 は 0(どれも $O(NK)$ DP と一致)
- 実行時間(実測): $N=10^5,K=500$ で 0.51 s、$N=10^5,K=5\times10^4$ で 0.54 s、$N=2\times10^5,K=1000$ で 1.03 s。$N=20000,K=1000$ では本解 0.17 s・$O(NK)$ DP 5.06 s で答えは一致。$N=3000$ では $K=10$・$K=300$ とも一致
- 扱っていない範囲: $g(k)$ の凸性の証明、ちょうど $K$ 回、IOI 2016 "Aliens" の原典確認、優先度付きキューによる別解の検証