CODE FESTIVAL 2016 Relay: K - 木の問題 / Problem on Tree

,

http://cf16-relay-open.contest.atcoder.jp/tasks/relay_k

darinflarさんとsigmaさんにほぼ任せていた。木DPであり話を聞く限り通りそうだったがWAが出て通らず。 後から解きなおして(WAを出して)気付いたが本番での自分の理解には抜けがあって、親方向の葉を見るのを忘れていた。

solution

木DP。各部分木ごとに、その部分木を経由する場合の$v$に含められる頂点の数(つまり葉の数)と、その部分木中の頂点を$v$の端点として用いた場合の$v$に含められる頂点の数を求める。$O(N)$。

答えとなる値の出し方には注意が必要である。 各頂点について、その頂点を使うとした場合はその子の中で(親方向は親が考えてくれるので無視してよい)端点として使ったときに大きい部分木を$2$つ選んで足せばよい。 一方その頂点を使わないとした場合は問題で、端点は両方その部分木から選ぶ(そうしたときの差分が大きい$2$つを選ぶ)としてよいが、親方向にある葉も使ってよいのでこれを足さねばならない。

implementation

#include <iostream>
#include <vector>
#include <algorithm>
#include <functional>
#define repeat(i,n) for (int i = 0; (i) < (n); ++(i))
using namespace std;
template <class T> void setmax(T & a, T const & b) { if (a < b) a = b; }
int main() {
    // input
    int n; cin >> n;
    vector<vector<int> > g(n);
    repeat (i,n-1) {
        int p, q; cin >> p >> q; -- p; -- q;
        g[p].push_back(q);
        g[q].push_back(p);
    }
    // compute
    int total_leaves = 0;
    repeat (i,n) if (g[i].size() == 1) total_leaves += 1;
    int ans = 0;
    vector<int> middle(n), terminal(n); // dp on tree
    function<void (int, int)> go = [&](int i, int parent) {
        // prepare
        int sum_middle = 0;
        vector<int> terminals;
        vector<int> diff;
        for (int j : g[i]) if (j != parent) {
            go(j, i);
            sum_middle += middle[j];
            terminals.push_back(terminal[j]);
            diff.push_back(terminal[j] - middle[j]);
        }
        diff.push_back(0);
        diff.push_back(0);
        sort(diff.rbegin(), diff.rend());
        terminals.push_back(0);
        terminals.push_back(0);
        sort(terminals.rbegin(), terminals.rend());
        // result
        middle[i] = max(1, sum_middle);
        terminal[i] = max(1 + terminals[0], sum_middle + diff[0]);
        setmax(ans, 1 + terminals[0] + terminals[1]);
        setmax(ans, total_leaves + diff[0] + diff[1]);
    };
    go(0, -1);
    // output
    cout << ans << endl;
    return 0;
}