こういうのを素早く解けるようになりたい

問題概要

$N$ 頂点 $M$ 辺の重み付き無向グラフがある。各頂点 $v$ には

そこで売ってるパンを買ったときの美味しさ $P_{v}$
そこで売ってるジャムの種類 $C_{v}$
そこで売ってるジャムを買ったときの美味しさ $J_{v}$

という 3 つの値がある。ジャムの種類は $1, 2, \dots, K$ のいずれかであることが保証されている。各 $k = 1, 2, \dots, K$ に対して、獲得する幸福度の最大値を求めよ。初期状態の幸福度を 0 として

頂点 1 を出発して
ある頂点 $u$ を訪れる (幸福度が、たどった辺の重みの和だけ減少)
頂点 $u$ でパンを買う (幸福度が、 $P_{u}$ だけ増加)
売ってるジャムの種類が $k$ であるような、ある頂点 $v$ を訪れる (幸福度が、たどった辺の重みの和だけ減少)
頂点 $v$ でジャムを買う (幸福度が、 $J_{v}$ だけ増加)
頂点 1 に戻る (幸福度が、たどった辺の重みの和だけ減少)

という操作を行う。

制約

$N \le 10^{5}$
$M \le 2 \times 10^{5}$

考えたこと

とりあえず、ジャムを買う頂点については決め打ちして、 $N$ 通りすべて試すことにする。このとき、パンを買う頂点もすべて決め打ちしてしまうと TLE しそうである。

そこで、次の値を管理する (幸福度を符号反転してコストと読み替えている)

dp1[ v ] := 頂点 1 から出発して、途中でパンを買わずに、頂点 v まで訪れるときの最小コスト
dp2[ v ] := 頂点 1 から出発して、途中どこかでパンを買った場合の、頂点 v まで訪れるときの最小コスト

つまり、頂点を「パンをまだ買ってない」「パンをすでに買った」という情報を持たせて倍加する。こういうの、古の競プロ界では拡張 Dijkstra 法などと呼ばれたりしていた。現在はそのネーミングはよくないよね、ってことになっている。

それはさておき、このままだと「パンを買う」という行為を表す遷移の重みが負になるので Dijkstra 法を使えなくなってしまう。Dijkstra 法が使えないと計算量的にやばそう。そこで、すべてのパンに一律に下駄を履かせて非負になるようにする。パンを買う個数はかならず 1 個なので、下駄を履かせて問題ない。

そして各頂点 v に対して、

J[ v ] - dp1[ v ] - (dp2[ v ] - GETA)

が v でジャムを買う場合の幸福度の最大値となる。

コード

計算量は $O(M \log N)$ となる。

#include <bits/stdc++.h>
using namespace std;
template<class T> inline bool chmax(T& a, T b) { if (a < b) { a = b; return 1; } return 0; }
template<class T> inline bool chmin(T& a, T b) { if (a > b) { a = b; return 1; } return 0; }

using Edge = pair<long long, int>;
using Graph = vector<vector<Edge>>;

int main() {
    int N, M, K;
    cin >> N >> M >> K;
    vector<long long> P(N), C(N), J(N);
    for (int i = 0; i < N; ++i) cin >> P[i];
    for (int i = 0; i < N; ++i) cin >> C[i];
    for (int i = 0; i < N; ++i) cin >> J[i];
    Graph G(N);
    for (int i = 0; i < M; ++i) {
        int u, v;
        long long t;
        cin >> u >> v >> t;
        --u, --v;
        G[u].emplace_back(v, t);
        G[v].emplace_back(u, t);
    }

    const long long INF = 1LL<<60;
    const long long GETA = 1LL<<30;
    vector<long long> dp1(N, INF), dp2(N, INF); // notuse pan, use pan
    using pint = pair<int,int>; // vertex, (notuse pan: 0, use pan: 1)
    priority_queue<pair<long long, pint>,
                   vector<pair<long long, pint>>,
                   greater<pair<long long, pint>>> que;
    dp1[0] = 0;
    que.push({0, pint(0, 0)});
    while (!que.empty()) {
        auto p = que.top(); que.pop();
        long long cur = p.first;
        int v = p.second.first, parity = p.second.second;

        if (parity == 0 && cur > dp1[v]) continue;
        if (parity == 1 && cur > dp2[v]) continue;

        if (parity == 0) {
            if (chmin(dp2[v], dp1[v] - P[v] + GETA)) {
                que.push({dp2[v], pint(v, 1)});
            }
            for (auto e : G[v]) {
                int nv = e.first, w = e.second;
                if (chmin(dp1[nv], dp1[v] + w)) {
                    que.push({dp1[nv], pint(nv, 0)});
                }
            }
        }
        else {
            for (auto e : G[v]) {
                int nv = e.first, w = e.second;
                if (chmin(dp2[nv], dp2[v] + w)) {
                    que.push({dp2[nv], pint(nv, 1)});
                }
            }
        }
    }

    vector<long long> res(K, -INF);
    for (int v = 0; v < N; ++v) {
        chmax(res[C[v]-1], J[v] - (dp1[v] + dp2[v] - GETA));
    }
    for (auto v : res) cout << v << endl;
}

けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AOJ 2956 ジャム (Jam) (HUPC 2019 day2-E)

問題概要

制約

考えたこと

コード