CodeChef June Challenge 2018: Archi and Tree

CodeChef June Challenge 2018の問題: Archi and Tree (Code: ARCTR)
問題ページ: https://www.codechef.com/JUNE18A/problems/ARCTR

問題概要

各辺がそれぞれの長さを持つ、個の頂点を含む木がある。
番目の辺はとを繋ぐ、長さの辺である。

この木の上に個の流れ星が流れる。
具体的に、番目の流れ星は時間に頂点からスタートし、速度でまっすぐに向かい、頂点に到達したら消える。

各頂点で最初にいずれかの流れ星が観測できる最小の時刻を出力せよ。観測できない頂点については-1を出力せよ。

制約

1つのテストケースに含まれるケース数:



入力されるグラフは木
全てのケースのの総和はを超えない
全てのケースのの総和はを超えない

解法

Heavy-Light Decomposition + Li Chao Tree で解いた。

例えば、下の図のようなケースであれば、

下の図のようにHeavy-Light Decompositionを用いてheavy-pathに沿って木をいくつかの木に分割し、

下の図のように分割された木の中で本の線分の最小値を計算すればよい。

HL分解された各木に線分を追加する処理は、頂点u, vが含まれたそれぞれの分割木から、頂点u, vのLCAの頂点wが含まれた分割木まで1つずつ親の木に遡りながら線分を追加する処理を行うことで達成できる。HL分解した木は、高さがになるため、この線分追加操作は本の線分に対し回になる。

ある点における、これらの線分が取る最小値はLi Chao Treeを用いることで計算できる。
Li Chao Treeでは、個の点(x座標)について計算を行う場合、1本の直線追加処理を、ある点xにおける最大値・最小値の計算を1回で計算できる。このLi Chao Treeで更に、線分が含まれる範囲のみに直線を追加するように更新することで、ある点における各線分のうちの最大値・最小値を計算できる。
参考: Convex hull trick and Li Chao tree - Competitive Programming Algorithms

実装

計算量は

HL分解を実装したつもりが、heavy-pathに沿ってないため実装できてない()。
雑な実装してる。

提出コード(Python3): https://www.codechef.com/viewsolution/18771280

import sys
sys.setrecursionlimit(10**6)
T = int(input())
for _ in range(T):
    N = int(input())
    G = [[] for i in range(N)]
    for i in range(N-1):
        u, v, w = map(int, input().split())
        G[u-1].append((v-1, w))
        G[v-1].append((u-1, w))
 
    # Heavy-Light Decomposition
    # DFSで探索しながら、木をいくつかの直線木に分割
    # lg[i]: 頂点iが属する分割木番号
    # li[i]: 頂点iの、分割木におけるインデックス
    # ld[i]: 頂点iの、分割木における左からの距離
    # grp[i]: 分割木iに含まれる頂点リスト
    # gprv[i]: 分割木番号iが繋がる親の分割木の情報 (gp, gi, gd)
    #    gp: 親の分割木番号
    #    gi: 分割木の中で繋がっている頂点インデックス
    #    gd: 分割木の中で繋がっている頂点の左からの距離
    # gdep[i]: 分割木iのrootの分割木からの深さ(距離)
    # gdist[i]: 分割木iに含まれる、各頂点の左からの距離
    lg = [0]*N; li = [0]*N; ld = [0]*N
    prv = [0]*N
    grp = []; gprv = []; gdep = []; gdist = []
    def hld(v, p, d, gp, gi, gd):
        L = []; D = []
        g = len(grp)
        grp.append(L)
        gdist.append(D)
        gprv.append((gp, gi, gd))
        gdep.append(d)
        di = i = 0
        while v is not None:
            lg[v] = g
            li[v] = i
            ld[v] = di
            D.append(di)
            L.append(v)
            nxt = None
            for w, co in G[v]:
                if w == p:
                    continue
                if nxt is None:
                    # vから最初に見つけた頂点を、現在の分割木の次の頂点とする
                    # 本来はheavy-pathにしなければならない
                    nxt = w
                    di += co
                else:
                    # 2つめ以降は、新たな分割木に含ませる
                    hld(w, v, d+1, g, i, co)
            p, v = v, nxt
            i += 1
    hld(0, -1, 0, -1, -1, 0)
 

    # Li Chao Tree

    # x=xにおける線分lineの値を計算
    def f(line, x):
        # 誤差を回避するため、r, p, qを保持し、f(x) = (r*x + q)/pと計算
        r, p, q = line
        if r == 1:
            return (x + q)/p
        return (-x + q)/p
 
    def _add_line(line, a, b, data, xs, k, l, r):
        if r <= a or b <= l:
            return
        m = (l + r) // 2
        if not (a <= l and r <= b):
            # 線分が含まれる区間まで移動
            _add_line(line, a, b, data, xs, 2*k+1, l, m)
            _add_line(line, a, b, data, xs, 2*k+2, m, r)
            return
        # 以下は線分が含まれる区間[l, r]での処理
        # 区間[l, r]に線分が存在しなければ、ここに線分追加して終了
        if data[k] is None:
            data[k] = line
            return
        lx = xs[l]
        mx = xs[m]
        rx = xs[r-1]
        r0, p0, q0 = line; r1, p1, q1 = data[k]
        left = (f(line, lx) < f(data[k], lx))
        mid = (f(line, mx) < f(data[k], mx))
        right = (f(line, rx) < f(data[k], rx))
        # 元の区間[l, r]の線分より全て値が小さければ、入れ替えて終了
        if left and right:
            data[k] = line
            return
        # 元の区間[l, r]の線分より全て値が大きければ、何もせず終了
        if not left and not right:
            return
        # 元の区間[l, r]の線分よりも、xs[m]における値が小さければ、ここで入れ替える
        if mid:
            data[k], line = line, data[k]
        # 今持っている線分が、区間[l, r]の線分よりも小さくなれる方の区間をさらに探索
        if left != mid:
            _add_line(line, a, b, data, xs, 2*k+1, l, m)
        else:
            _add_line(line, a, b, data, xs, 2*k+2, m, r)
 
    # 区間[xs[a], xs[b]]の線分lineを追加する
    def add_line(line, a, b, i):
        return _add_line(line, a, b+1, ldata[i], gdist[i], 0, 0, NN[i])
 
    # xs[k] (= x)におけるM本の線分がとる最小値を計算
    def query(k, x, data, n):
        k += n-1
        s = 1e30
        while k >= 0:
            if data[k]:
                s = min(s, f(data[k], x))
            k = (k - 1) // 2
        return s
 
    LL = len(grp)
    NN = [2**(len(e)-1).bit_length() for e in grp]
    ldata = [[None]*(2*n+1) for n in NN]
 
    M = int(input())
    for i in range(M):
        u, v, t, s = map(int, input().split()); u -= 1; v -= 1
        ug = lg[u]; vg = lg[v]
        if lg[u] == lg[v]:
            # 同じ分割木にu, vが存在
            ui = li[u]; vi = li[v]
            ud = ld[u]; vd = ld[v]
            if ui <= vi:
                add_line((1, s, t*s-ud), ui, vi, ug)
            else:
                add_line((-1, s, t*s+ud), vi, ui, ug)
        else:
            # 異なる分割木にu, vが存在

            # 同じ高さになるように親の木に遡る
            udp = gdep[ug]; vdp = gdep[vg]
            ud = ld[u]; ui = li[u]
            d = 0
            while udp > vdp:
                # 根頂点に近づくにつれて時間が増大する線分を追加
                # w = LCA(u, v)とすると、u -> w -> vのうちのu -> wの部分
                ud = gdist[ug][ui]
                d += ud
                add_line((-1, s, t*s+d), 0, ui, ug)
                ug, ui, ud = gprv[ug]
                d += ud
                udp -= 1
            st = []
            vd = ld[v]; vi = li[v]
            while udp < vdp:
                # w = LCA(u, v)とすると、u -> w -> vのうちのw -> vの部分は、uからの距離が
                # 分かってから計算を行うためスタックする
                vd = gdist[vg][vi]
                st.append((vd, vg, vi))
                vg, vi, vd = gprv[vg]
                vdp -= 1
 
            # 同じ木に到達するまで１つずつ親の木に遡る
            while ug != vg:
                ud = gdist[ug][ui]
                d += ud
                add_line((-1, s, t*s+d), 0, ui, ug)
                ug, ui, ud = gprv[ug]
                d += ud
                udp -= 1
 
                vd = gdist[vg][vi]
                st.append((vd, vg, vi))
                vg, vi, vd = gprv[vg]
                vdp -= 1
 
            # 同じ木に到達した時の線分追加処理
            ud = gdist[ug][ui]; vd = gdist[vg][vi]
            if ui <= vi:
                add_line((1, s, t*s+(d-ud)), ui, vi, ug)
            else:
                add_line((-1, s, t*s+(d+ud)), vi, ui, ug)
            d += abs(ud - vd)

            # 根頂点に近づくにつれて時間が減少する線分を追加
            # w = LCA(u, v)とすると、u -> w -> vのうちのw -> vの部分
            while st:
                e, vg, vi = st.pop()
                vd = gprv[vg][2]
                d += vd
                add_line((1, s, t*s+d), 0, vi, vg)
                d += e
 
    ans = []
    for i in range(N):
        # 各頂点について、最小値を計算
        g = lg[i]
        j = li[i]
        d = ld[i]
 
        y = query(j, d, ldata[g], NN[g])
        if y >= 1e29:
            ans.append("-1")
        else:
            ans.append("%.08f" % y)
    print(*ans, sep='\n')