AtCoder Grand Contest 010: C - Cleaning

,

http://agc010.contest.atcoder.jp/tasks/agc010_c

Aを書いた後それをBかCが書けたら投げようと思ったら、BもCも最後まで分からなかったので無提出rating変化なし。 あまり良くないとは思うのだけれど、短期的なratingを考えると十分選択肢に入る戦略なのでこうなってしまう。

木DPだとは思っていたが、各部分木を全てちょうど葉と同一視できるのに気付かず、その頂点が要求するパスの数の加減$L_i$と上限$R_i$で木DPしようとしていた。

solution

木DP。部分木は葉と同一視できるので根が$A_i = 0$な葉と見做せるかを答える。$O(N)$。

ある葉でない部分木$i$を葉と同一視するとして持つ石の数$A’_i$を考える。 子は全て葉としてよい。 子の持つ石の総和$\mathrm{sum} = \sum_{j \in \mathrm{children}(i)} A_j$とする。 この部分木の中で閉じたパスの数を$b$、外へ出ていくパスの数を$c$とすると、$A_i = b + c$かつ$\mathrm{sum} = 2b + c$である。 これは連立方程式として解けて$b = \mathrm{sum} - A_i$。一意に定まることに注意。 また$c = A_i - b$は求めたい$A’_i$と等しい。 ここでそのような$(b, c)$がパスの張り方として有効かどうかの確認が必要。 $0 \le b, 0 \le c$に加えて、内部で$b$本張れるかを見れば十分。 子で$A_j$が最も大きいものを$j$として、$A_j \ge \frac{\mathrm{sum}}{2}$であればこの子が制限をして$b \le \mathrm{sum} - A_j$である必要がある。 そうでないとすると、単に$\mathrm{sum}$が制限をして$b \le \frac{\mathrm{sum}}{2}$を確認すればよい。

implementation

#include <iostream>
#include <vector>
#include <algorithm>
#include <functional>
#include <cassert>
#define repeat(i,n) for (int i = 0; (i) < int(n); ++(i))
#define whole(f,x,...) ([&](decltype((x)) whole) { return (f)(begin(whole), end(whole), ## __VA_ARGS__); })(x)
using ll = long long;
using namespace std;
bool solve(int n, vector<int> & a, vector<vector<int> > & g) {
    vector<int> dp(n);
    auto is_leaf = [&](int i) { return g[i].size() == 1; };
    function<bool (int, int)> go = [&](int i, int parent) {
        if (is_leaf(i)) {
            dp[i] = a[i];
        } else {
            ll sum_dp = 0;
            for (int j : g[i]) if (j != parent) {
                if (not go(j, i)) return false;
                sum_dp += dp[j];
            }
            ll b = sum_dp - a[i];
            ll c = a[i] - b;
            if (b < 0 or c < 0) return false;
            int argmax_dp = *whole(max_element, g[i], [&](int j, int k) { return make_pair(j != i, dp[j]) < make_pair(k != i, dp[k]); });
            ll max_dp = dp[argmax_dp];
            if (max_dp > sum_dp / 2) {
                if (sum_dp - max_dp < b) return false;
            } else {
                if (sum_dp / 2 < b) return false;
            }
            dp[i] = c;
        }
        return true;
    };
    assert (n >= 2);
    if (n == 2) {
        return a[0] == a[1];
    } else {
        int i = 0;
        while (is_leaf(i)) ++ i;
        if (not go(i, -1)) return false;
        return dp[i] == 0;
    }
}
int main() {
    int n; cin >> n;
    vector<int> a(n); repeat (i,n) cin >> a[i];
    vector<vector<int> > g(n);
    repeat (i,n-1) {
        int x, y; cin >> x >> y; -- x; -- y;
        g[x].push_back(y);
        g[y].push_back(x);
    }
    cout << (solve(n, a, g) ? "YES" : "NO") << endl;
    return 0;
}