CHT (Convex Hull Trick) のpython実装を書いた

AtCoder のEducational DP contest-Z Frogを 解いていて必要だったので整備しました。

自分の実装 : Submission #9866424 - Educational DP Contest

CHTについてはググればたくさんの記事がある。以下の２つの関数がある。

1. insert 関数 : ax+bを追加する。単調減少である順にaを追加する場合と一般の場合で計算量が変わる。
2. find関数 : あるxに対する最小値を求める。 xが単調増加になる順にfind関数を作用させる場合と一般の場合で計算量が変わる。

実装

lines = []
def fi(i, x):
    a, b = lines[i]
    return a*x+b

def find(x):
    def f(i):
        return fi(i+1,x) > fi(i,x)
    mn, mx = -1, len(lines)-1
    idx = (mn+mx)//2
    while mx-mn>1:
        if f(idx):
            mx, idx = idx, (mn + idx)//2
            continue
        mn, idx = idx, (mx + idx)//2
    return fi(idx+1, x)

def find2(x):
    # x が単調増加
    I = find2.I
    u, v = lines[I]
    while I+1<len(lines):
        w, y = lines[I+1]
        if u*x+v < w*x+y:
            break
        u, v = w, y
        I += 1
    find2.I = I
    return u*x+v
find2.I = 0

def insert(a, b):
    pass

def insert2(a, b):
    # a が単調減少
    if not lines:
        lines.append((a,b))
        return
    (e, f) = lines[-1]
    while len(lines)-1:
        (c, d), (e, f) = (e, f), lines[-2]
        if (c-e)*(b-d) < (d-f)*(a-c):
            break
        lines.pop()
    lines.append((a, b))

N, = map(int, input().split())
As = map(int, input().split())
for j, a in enumerate(As):
    j+=1
    insert2(-2*j, a+j**2)
for i in range(N):
    print(find(i+1)+(i+1)**2)

insert関数、find関数はそれぞれ一般の場合だがinsertは未実装。

insert2関数、find2関数はそれぞれ傾きが単調減少の場合、xが単調増加の場合で高速に動作する。

linesがglobalな変数となっていてtuple(a, b) (a:傾き, b:切片)のリストになっていて、updateやinsertを実行すると内部が更新される。