CodeChef July Challenge 2018: Pizza Delivery - 日々ｄｒｄｒする人のメモ

CodeChef July Challenge 2018の問題: Pizza Delivery (Code: PDELIV)
問題ページ: https://www.codechef.com/JULY18A/problems/PDELIV

問題概要

一直線上に ${n}$ 個のピザ屋と ${m}$ 人のお客がいて、 ${i}$ 番目のピザ屋は ${s_i}$ 、 ${i}$ 番目のお客は ${c_i}$ にいる。

${i}$ 番目のピザ屋で1枚注文する場合、ベースの値段 ${p_i}$ に加えて、ピザ屋の位置 ${x_1}$ からお客の位置 ${x_2}$ まで配達するのに追加料金 ${(x_1 - x_2)^2}$ がかかる。

また、各お客 ${i}$ には注文したくないピザ屋が ${k_i}$ 個存在する。
この状況で、 ${m}$ 人のお客がピザ1枚注文するための、最小料金はそれぞれいくらか？

制約

${1 \le m, n \le 2 \times 10^5}$
${0 \le s_i, c_i \le 10^9}$
${1 \le p_i \le 10^9}$
${0 \le k_i \le n - 1}$
${0 \le \sum_{i=1}^m k_i \le 4 \times 10^5}$
${s_1, s_2, ..., s_n}$ はそれぞれ異なる座標

解法

Li Chao Tree上で各お客について、線分の最小値を計算すると解ける。

まず、お客 ${i}$ に対して解くべき式は以下のように変形できる。( ${S_i}$ はお客 ${i}$ が注文できるピザ屋のインデックス集合)

$\min_{j \in S_i} ( p_j + (s_j - c_i)^2) = c_i^2 + \min_{j \in S_i} (-2c_is_j + s_j^2 + p_j)$

${f_j(x) = -2 s_j x + s_j^2 + p_j}$ とおけば、Convex Hull Trickで解ける問題になる。

この問題では、各お客について注文したくないピザ屋が存在するため、これをうまく扱う必要がある。
これは下の図のように、各ピザ屋 ${j}$ の直線 ${f_j}$ について、そのピザ屋を注文したくないお客 ${i}$ の座標だけその直線を追加しない(つまり、注文できないお客 ${i}$ について ${x < c_i, c_i < x}$ となるように直線を切った線分にする)ことで解決できる。
f:id:smijake3:20180716233119p:plain

そして、これらの線分の最小値はLi Chao Tree上で計算することができる。
smijake3.hatenablog.com

追加する線分は、各ピザ屋について、注文できないお客の人数だけ行うことになるため、全体で高々 ${(\sum_{i=1}^m k_i + n)}$ 本になる。

しかし問題の制約上、 ${i \neq j}$ について ${c_i = c_j}$ となるケースが存在するが、同じ座標で同じピザ屋を注文できるかが異なるとダメなので、複数のお客を同じ座標で扱うことを回避する必要がある。
そこで、Li Chao Tree上で同じ座標の複数のお客を区別して扱い、それぞれに対し線分を追加するしないを決めるようにする。つまり ${x_1 \le x_2 \le ... \le x_m}$ となる ${m}$ 個の座標をLi Chao Treeで管理する。このように同じ座標が複数存在する場合でも、Li Chao Tree上で正しく扱える。

最後に、全ての線分を追加したら、各お客について料金の最小値を計算すればよい。

実装

${K = \sum_{i=1}^m k_i}$ とすると、計算量は ${O((K + n) \log^2 m + m \log m)}$

提出コード (PyPy2): https://www.codechef.com/viewsolution/19144188

N, M = map(int, raw_input().split())
 
# Li Chao Tree
N0 = 2**(M-1).bit_length()
data = [None]*(2*N0+1)
 
def f(line, x):
    p, q = line
    return p*x + q
 
def _add_line(line, k, l, r):
    m = (l + r) // 2
    if data[k] is None:
        data[k] = line
        return
    lx = X[l]; mx = X[m]; rx = X[r-1]
    left = (f(line, lx) < f(data[k], lx))
    mid = (f(line, mx) < f(data[k], mx))
    right = (f(line, rx) < f(data[k], rx))
    if left and right:
        data[k] = line
        return
    if not left and not right:
        return
    if mid:
        data[k], line = line, data[k]
    if left != mid:
        _add_line(line, 2*k+1, l, m)
    else:
        _add_line(line, 2*k+2, m, r)
 
def add_line(line, a, b):
    L = a + N0; R = b + N0
    a0 = a; b0 = b
    sz = 1
    while L < R:
        if R & 1:
            R -= 1
            b0 -= sz
            _add_line(line, R-1, b0, b0+sz)
        if L & 1:
            _add_line(line, L-1, a0, a0+sz)
            L += 1
            a0 += sz
        L >>= 1; R >>= 1
        sz <<= 1
 
def query(k):
    x = X[k]
    k += N0-1
    s = 1e30
    while k >= 0:
        if data[k]:
            s = min(s, f(data[k], x))
        k = (k - 1) // 2
    return s
 
# F[i]: ピザ屋iで注文しないお客リスト
F = [[] for i in xrange(N)]

C = []
C0 = [0]*M
X = [10**10]*N0
PM = [map(int, raw_input().split()) for i in xrange(N)]
for i in xrange(M):
    pp = map(int, raw_input().split())
    # c_i
    c = pp[0]
    # 各d_ij
    for e in pp[2:]:
        F[e-1].append(i)
    C.append((c, i))

    C0[i] = c

# c_i をソートする (元のインデックスを取得できるようにしておく)
MP = [None]*M
C.sort()
for i, (c, j) in enumerate(C):
    X[i] = c  # i番目の座標はc
    MP[j] = i # j番目のお客はソートしてi番目
 
# 各ピザ屋ごとに、注文しないお客を飛ばすように直線を切りながら線分追加していく
for i, (s, p) in enumerate(PM):
    # お客リストを座標に変換してソート
    Fi = map(MP.__getitem__, F[i])
    Fi.sort()
    Fi.append(M)

    # f(x) = -2*s*x + p+s**2 を追加
    prv = -1
    line = (-2*s, p+s**2)
    for e in Fi:
        if prv+1 < e:
            add_line(line, prv+1, e)
        prv = e
 
# 各お客ごとに最小コストを計算
ans = []
for i in xrange(M):
    ans.append(query(MP[i]) + C0[i]**2)
print("\n".join(map(str, ans)))