AOJ 2725 Live Programming - すいバカ日誌

問題文

http://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=2725

問題概要

催し物を時間 $T$ だけ開催することになった．
そこで歌う歌の候補が $N$ 個あり，それぞれの歌はパラメーター $t[i], p[i], f[i]$ を持つ．$t[i]$ は歌の長さ，$p[i]$ は歌の満足度，$f[i]$ は歌の特徴量を表します．これらの歌からいくつか選んで，歌うことにする．
この時，以下の制約を満たすような歌の列 $\{s_1, ..., s_k\}$ で，満足度が最大を取るときの満足度を求めよ．

歌う歌の時間の総和が $T$ 以下である．
最初に歌う歌の満足度は $p[s_0]$ とする．
以降に歌う歌の満足度は，$p[s_i] - (f[s_{i-1}] - f[s_i])^2$ である．$s_{i-1}$ は，直前に歌った歌である．

・制約
1 <= N <= 4000
1 <= T <= 4000

解法

DPっぽい感じがするのでDPでやるんですが，その前に考えるべきことがあります．

今回問題になるのは，「どの歌の集合を選ぶか」と「歌をどの順番で並べるか」の2つあり，このままではしんどいです．
では，仮に歌の集合を決め打ちした場合，それらをどう並べるのがいいか考えてみましょう．これは明らかに，f[i] の大きさでソートされた順番に歌うのが最適です．

つまり，あらかじめ f[i] の値で歌をソートしてしまえば，以下のDPでできることは明白です．
$dp[i][j] := i$ 番目の歌を最後に歌うとして，総時間が $j$ であるときの最大満足度
このDPの遷移は以下のとおりです．
$\displaystyle dp[i][j] = \max_{k < i}\{dp[k][j - t[i]] + p[i] - (f[k] - f[i])^2\}$
式を整理して，さらに（別に置き換えなくてもできますが）max を min に置き換えると以下の式になります．
$\displaystyle dp[i][j] = -\min_{k < i}\{-dp[k][j - t[i]] + (f[k])^2 - 2f[k] \times f[i]\} + p[i] - (f[i])^2$
これはConvexHullTrickチャンスなので高速化でき，計算量 O(NT) になります．

注意すべき点は j のループは大きい方から小さい方に回すことです．これは普通のナップザックと同じです．
こういう細かいミスをなくすには，先に自明なDPを書いてから，あとで高速化するとよいですね．（慣れているなら別）

ソースコード

#include <bits/stdc++.h>
using namespace std;
 
using ll = long long;
 
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];
    }
 
    bool empty() const {
        return a.size() == 0;
    }
 
private:
    std::deque<T> a, b; // functions
};
 
int main() {
    int N, T;
    cin >> N >> T;
    // (f, t, p)
    vector<tuple<ll, ll, ll>> songs;
    for(int i = 0; i < N; ++i) {
        ll t, p, f;
        cin >> t >> p >> f;
        if(t > T) {
            continue;
        }
        songs.emplace_back(f, t, p);
    }
    N = songs.size();
    sort(songs.begin(), songs.end());
 
    vector<vector<ll>> dp(N + 1, vector<ll>(T + 1));
    vector<convex_hull_trick> cht(T + 1);
    ll res = 0;
    for(int i = 0; i < N; ++i) {
        ll fi, ti, pi;
        tie(fi, ti, pi) = songs[i];
        for(int j = T; j >= ti; --j) {
            if(cht[j - ti].empty()) {
                continue;
            }
            dp[i][j] = -cht[j - ti].min(fi) - fi * fi + pi;
            res = max(res, dp[i][j]);
            cht[j].add_f(-2 * fi, -dp[i][j] + fi * fi);
        }
        dp[i][ti] = max(dp[i][ti], pi);
        res = max(res, dp[i][ti]);
        cht[ti].add_f(-2 * fi, -dp[i][ti] + fi * fi);
    }
    cout << res << endl;
}