Codeforces Round #185 (Div. 1) B. Cats Transport

問題文

http://codeforces.com/problemset/problem/311/B

問題概要

N個の地点があり，順に 0 から N-1 とする．地点 i と i + 1 の間の距離は d[i] である．
また，猫が m 匹いて，猫 j は時刻 t[j] に地点 h[j] にあらわれるものとする．それぞれの猫は，餌をもらえるまで出現してからその地点で待機し続ける．
今，猫に餌を上げる人が p 人いて，それぞれ地点 0 から N-1 の方向に単位時間あたり1だけ移動する．（ただし初期位置は自由．つまり負の時刻に地点 0 を出発するような人がいても良い．これは制約に書いてないが…）
移動中に餌を待っている猫と遭遇した場合，その猫に必ず餌をやる．ただし，餌をやる時間は無視できるものとする．
では，全ての猫に餌をやるように人を配置した時に，それぞれの猫が餌を待つ時間の総和を最小化せよ．

・制約
2 <= N <= 2 * 10^5
1 <= M <= 10^5
1 <= p <= 100

解法

まず，今のままでは位置と時間という軸が2つあり面倒なので，これをひとつにまとめます．人は単位時間に1だけすすむので，時刻 t[j] に地点 h[j] に現れるという情報を a[j] := t[j] - pos[j] とまとめてしまいます（pos[j] は地点 h[j] の座標）．a[j] は，時刻 a[j] に人を地点 0 に配置すれば，その人は猫 j が現れるよりあとに通過する（つまり餌をあげることができる）ことを意味します．
こうして得られた a[j] を昇順にソートしておきます．
時刻 a[j] に人を地点 0 に配置した場合，それ以前の直近で人が配置された時刻を a[k] をすると，k + 1 <= j となるような全ての猫について餌をやることになります．

ここからDPパートです．
p が小さいので，以下のような dp が自然に思い浮かびます．
dp[i][j] := i 人目まで使い，かつ i 人目を時刻 a[j] に出発させた時の，j 以前の猫の待ち時間の総和の最小値
このdpの更新は以下のようになります．
$\displaystyle dp[i][j] = \min_{k < j} \{dp[i - 1][k] + \sum_{l=k+1}^{j}(a[j] - a[l])\}$
式を変形すると，以下のようにできます．
$\displaystyle dp[i][j] = \min_{k < j} \{dp[i - 1][k] - k \times a[j] - \sum_{l=k+1}^{j}a[l]\} + j \times a[j]$
さらに累積和 $\displaystyle S[j] = \sum_{l=0}^{j-1} a[l]$を用いて整理すると
$\displaystyle dp[i][j] = \min_{k < j} \{dp[i - 1][k] - k \times a[j] + S[k+1] \} + j \times a[j] - S[j+1]$
このままだと O(pM^2) ですが，この形のDPは ConvexHullTrick という方法で高速化できます．（この方法自体の解説は省略します．）
よって O(pM) でこの問題を解くことができます．

ソースコード

#include <bits/stdc++.h>
using namespace std;

class convex_hull_trick {
    using T = long long;

    bool check(int f1, int f2, T aa, T bb) {
        return (a[f2] - a[f1]) * (bb - b[f2]) >= (b[f2] - b[f1]) * (aa - a[f2]);
    }
    bool check(int f1, int f2, int f3) {
        return (a[f2] - a[f1]) * (b[f3] - b[f2]) >= (b[f2] - b[f1]) * (a[f3] - a[f2]);
    }

    T f(int idx, T x) {
        return a[idx] * x + b[idx];
    }

public:
    convex_hull_trick() = default;
    // a : sorted sequence
    convex_hull_trick(std::vector<T> const& a_, std::vector<T> const& b_) {
        assert(a_.size() == b_.size());
        for(int i = 0; i < a_.size(); ++i) {
            add_f(a_[i], b_[i]);
        }
    }

    void add_f(T aa, T bb) {
        while(a.size() >= 2 && check(a.size() - 2, a.size() - 1, aa, bb)) {
            a.pop_back();
            b.pop_back();
        }
        a.push_back(aa);
        b.push_back(bb);
    }
    T min(T x) {
        assert(a.size() > 0);
        while(a.size() >= 2 && f(0, x) >= f(1, x)) {
            a.pop_front();
            b.pop_front();
        }
        return a[0] * x + b[0];
    }

private:
    std::deque<T> a, b; // functions
};

using namespace std;

using ll = long long;

constexpr ll INF = 2e18;

int main() {
    ll n, m, p;
    cin >> n >> m >> p;
    vector<ll> d(n - 1);
    for(int i = 0; i < n - 1; ++i) {
        cin >> d[i];
    }
    vector<ll> pos(n);
    for(int i = 0; i < n - 1; ++i) {
        pos[i + 1] = pos[i] + d[i];
    }
    vector<ll> x(m);
    for(int i = 0; i < m; ++i) {
        ll h, t;
        cin >> h >> t;
        x[i] = t - pos[h - 1];
    }
    sort(x.begin(), x.end());
    vector<ll> sum(m + 1);
    for(int i = 0; i < m; ++i) {
        sum[i + 1] = sum[i] + x[i];
    }

    vector<vector<ll>> dp(p + 1, vector<ll>(m, INF));
    vector<convex_hull_trick> cht(p + 1);
    for(int i = 0; i < m; ++i) {
        dp[1][i] = i * x[i] - sum[i];
        cht[1].add_f(-i, dp[1][i] + sum[i + 1]);
    }
    for(int i = 2; i <= p; ++i) {
        for(int j = 0; j < m; ++j) {
            if(j == 0) {
                dp[i][j] = 0;
            } else {
                dp[i][j] = cht[i - 1].min(x[j]) + j * x[j] - sum[j + 1];
            }
            cht[i].add_f(-j, dp[i][j] + sum[j + 1]);
        }
    }
    cout << dp[p][m - 1] << endl;
}

感想

ConvexHullTrickを初めて書いたので大変だった．かなり典型問題なので勉強になりました．