AOJ 2344 Multi Ending Story - すいバカ日誌

問題文

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

解法

木DPで解きます．
まず部分木に注目しましょう．
冷静に考えると，この部分木のすべての葉を回収するのには，以下の3通りだけしかありません．この部分木の根を $v$，左右の子をそれぞれ $l, r$ とします．

$v$ で quick save を使わず，それぞれの子の部分木でいい感じに quick save を用いて最適に回収する
$v$ で quick save した後，まず $l$ の部分木の葉をすべて回収（この時 quick save は用いない）して，その後 $r$ の部分木で quick save を用いつつ最適に回収する
2つ目のパターンで左右を入れ替えたパターン

そこで以下のように定めます．
$dp[v] := v$ が根の部分木にあるすべての葉を quick save をいい感じに用いて最適に回収したときのコスト
$dp_2[v] := v$が根の部分木にあるすべての葉を， $v$ で quick save をしたとき（子の部分木では一切 quick save しない）の回収のコスト．ただし真の根から $v$ までのコストは除くものとする．
$cnt[v] := v$ の部分木に含まれる葉の個数

1つ目のパターンは，$v$ で quick save しないので真の根から $v$ までのコスト $depth[v]$ が余計にかかります．なので，このときのコストは
$dp[l] + dp[r] + depth[v] + 2$
となります．

2つ目（3つ目も同様）のパターンでは，$l$ の部分木の回収に $dp_2[l]$ と，
$v$ から $l$ に $cnt[l]$ 回移動することになるのでその分だけのコストがまずかかります．その後は，$r$ にうつって最適に回収するので，そのコストは $dp[r] + 1$ です．よって総コストは
$dp_2[l] + cnt[l] + dp[r] + 1$
となります．

あとはこれらのうちで最小のものが $dp[v]$ の値となります．

ソースコード

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

using ll = long long;

// スタック拡張します

int N;
ll dp[500001], dp2[500001];
int leaf_cnt[500001];
int g[500001][2];

ll dfs(int v, int depth) {
    for(int i = 0; i < 2; ++i) {
        int to = g[v][i];
        if(to == N) {
            leaf_cnt[v]++;
        } else {
            dfs(to, depth + 1);
            leaf_cnt[v] += leaf_cnt[to];
            dp2[v] += dp2[to];
        }
    }
    dp2[v] += leaf_cnt[v];

    ll& res = dp[v];
    int l = g[v][0], r = g[v][1];
    dp[v] = min({dp2[l] + (leaf_cnt[l] == 0 ? 1 : leaf_cnt[l]) + 1 + dp[r],
                 dp[l] + dp2[r] + (leaf_cnt[r] == 0 ? 1 : leaf_cnt[r]) + 1,
                 dp[l] + depth + dp[r] + 2});
    return dp[v];
}

int main() {
    cin >> N;
    for(int i = 0; i < N - 1; ++i) {
        int a, b;
        cin >> a >> b;
        a -= (a != N);
        b -= (b != N);
        g[i][0] = a;
        g[i][1] = b;
    }
    cout << dfs(0, 0) << endl;    
}

感想

メモリの制約が厳しい上に，普通にスタックオーバーフローするので†スタック拡張†します．
DP部分は普通に難しかった……．