AOJ 2725 Live Programming

問題概要

催し物を時間 $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;
}