Learning Algorithms

アルゴリズムの勉強メモ

Atcoder Regular Contest 063 E. 木と整数 / Integers on a Tree

Atcoder Regular Contest 063 E. 木と整数 / Integers on a Tree

E - 木と整数 / Integers on a Tree

感想

ふつうに木の上で計算しました。

解法

まず頂点 $0$ を根として考えることにします。明らかな性質として、深さが同じような頂点に書かれた値すべての偶奇は一致していないといけません。最初にこれを判定しておきます。

そのあと、各部分木について、その子どもがとりうる値の区間(自分自身に数値が書かれていれば、その区間も加えて) $AND$ 区間をとればその頂点におけるとりうる区間が定められます。$AND$ 区間を求めるためには、すべての区間の $l$ の $max$ と、 $r$ の $min$ をとればよいです。

これを繰り返していった結果、頂点 $0$ に有効な区間が残っていれば、答えは $Yes$ で、根から順に適切に決めていくことができます。

実装
#include <algorithm>
#include <vector>
#include <functional>
using namespace std;

static const int INF = 0x3f3f3f3f;

void ex() {
        puts("No");
        exit(0);
}

int main() {
        int n;
        scanf("%d", &n);
        vector<vector<int>> g(n);
        for (int i = 0; i < n - 1; i ++) {
                int a, b;
                scanf("%d%d", &a, &b);
                a --, b --;
                g[a].push_back(b);
                g[b].push_back(a);
        }
        int k;
        scanf("%d", &k);
        vector<int> val(n, INF);
        while (k --) {
                int v, p;
                scanf("%d%d", &v, &p);
                v --;
                val[v] = p;
        }
        //check if even-odd prerequisites are met
        vector<int> depth(n);
        function<void (int, int, int)> get_depth = [&](int u, int prev, int d) {
                depth[u] = d;
                for (auto v : g[u]) if (v != prev) {
                        get_depth(v, u, d + 1);
                }
        };
        get_depth(0, -1, 0);
        vector<int> odd, even;
        for (int i = 0; i < n; i ++) {
                if (depth[i] & 1) {
                        if (val[i] != INF) odd.push_back(val[i]);
                } else {
                        if (val[i] != INF) even.push_back(val[i]);
                }
        }
        if (odd.size() > 0) {
                for (int i = 0; i < (int) odd.size(); i ++) {
                        if (odd[0] % 2 != odd[i] % 2) {
                                ex();
                        }
                }
        }
        if (even.size() > 0) {
                for (int i = 0; i < (int) even.size(); i ++) {
                        if (even[0] % 2 != even[i] % 2) {
                                ex();
                        }
                }
        }
        //compute
        vector<pair<int, int>> seg(n);
        function<void (int, int)> dfs = [&](int u, int prev) {
                vector<pair<int, int>> segs;
                for (auto v : g[u]) if (v != prev) {
                        dfs(v, u);
                        if (seg[v].first != INF) {
                                segs.push_back(make_pair(seg[v].first - 1, seg[v].second + 1));
                        }
                }
                if (segs.empty()) {
                        if (val[u] != INF) {
                                seg[u] = make_pair(val[u], val[u]);
                        } else {
                                seg[u] = make_pair(INF, INF);
                        }
                        return;
                }
                int l = 0, r = INF;
                for (int i = 0; i < segs.size(); i ++) {
                        l = max(l, segs[i].first);
                        r = min(r, segs[i].second);
                }
                if (val[u] != INF) {
                        if (val[u] < l || r < val[u]) {
                                ex();
                        } else {
                                l = val[u], r = val[u];
                        }
                }
                if (l > r) {
                        ex();
                }
                seg[u] = make_pair(l, r);
        };
        dfs(0, -1);
        //constructable
        printf("Yes\n");
        vector<int> ans(n, INF);
        function<void (int, int)> decide = [&](int u, int prev) {
                for (auto v : g[u]) if (v != prev) {
                        if (seg[v].first == INF) {
                                ans[v] = ans[u] + 1;
                        } else {
                                int l = seg[v].first, r = seg[v].second;
                                if (l <= ans[u] + 1 && ans[u] + 1 <= r) {
                                        ans[v] = ans[u] + 1;
                                } else {
                                        ans[v] = ans[u] - 1;
                                }
                        }
                        decide(v, u);
                }
        };
        ans[0] = seg[0].first;
        decide(0, -1);
        for (int i = 0; i < n; i ++) {
                printf("%d\n", ans[i]);
        }
        return 0;
}