Atcoder AGC #018 F. Two Trees

F: Two Trees - AtCoder Grand Contest 018 | AtCoder

頂点数が等しい根付き木が２つ与えられる。この木の各ノードに対して、値を割り当てることを考える。ただしノード番号が同じものは同じ値を割り当てる。任意の部分木（根を含む）について、それが含む各頂点に割り当てられた値の総和の絶対値が１に等しくなるような、値の割り当て方が存在するかどうか、そして存在するならばその一例を示すという問題。

半分自分で考え、半分解説に頼った。まず、条件を満たす方法によって、あるノードに値が割り当てられているならば、その値は、そのノードの子（子孫ではなく）の数の偶奇と一致していてはいけない。なぜならば、条件より、その子以下からなる根付き木について、その総和は＋１か−１のどちらかであるので、子ノードが一つ増えるということは、総和の偶奇が反転するということであるからである。さらにこのことより、各ノードの子の数の偶奇はどちらの木においても同じであることが必要である。もう一つ気づかなければならないことは、割り当てる値を必要以上に大きくしなくてよいことである。どこの部分木の和をとってもその絶対値が１を超えないようにすればよいので、総和に関しても、割り当てる値に関しても、-1, 0, 1以外の値を取らせなくてよい。このことから、各ノードに対して、その子の数を数えることで、その偶奇に応じて、0を割り当てるか、1 or -1を割り当てるかを決めることができる。1 or -1の部分を部分木内でマッチングさせて、あとでマッチングされたもの同士が1と-1になるようにうまく割り当てる。このとき、2つの木の両方で条件を満たすように割り当てなければならないので、このマッチングさせたもの同士が辺で結ばれた新しいグラフを用意し、そこで隣り合うもの同士を別の色で塗ることを考えれば、すべてうまく割り当てることができる。閉路ができることもあるが、長さが偶数なので、問題ない。

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

vector<int> g[2][101010];
vector<int> gg[202020];
vector<int> ans;
int used[101010];

int dfs(int v, int k, vector<vector<int>> &zero) {
        vector<int> matching; //ノードvを根とする部分木の中で、マッチングさせたいものをいれる
        if (!zero[k][v]) matching.push_back(v);
        for (int u : g[k][v]) {
                int get = dfs(u, k, zero);
                if (get == -1) continue;
                matching.push_back(get);
        }
        if (g[k][v].size() == 0) return (zero[k][v] ? -1 : v);
        while (matching.size() >= 2) { //２個ずつてきとうにマッチングさせて良い
                int s1 = matching.back();
                matching.pop_back();
                int s2 = matching.back();
                matching.pop_back();
                gg[s1].push_back(s2);
                gg[s2].push_back(s1);
        }
        if (matching.size() == 1) return matching[0]; //余った場合はマッチングさせたいノード番号を返す
        else return -1; //これ以上マッチングさせるものはない
}

void coloring(int v, int w) {
        used[v] = true;
        ans[v] = w;
        for (int u : gg[v]) if (!used[u]) {
                coloring(u, -w); //1と-1を交互に塗る
        }
}

int main() {
        int n, a;
        cin >> n;
        vector<int> root(2);
        for (int k = 0; k < 2; k ++) {
                for (int i = 0; i < n; i ++) {
                        cin >> a;
                        if (a == -1) { 
                                root[k] = i;
                        } else { 
                                a --;
                                g[k][a].push_back(i);
                        }
                }
        }
        vector<vector<int>> zero(2, vector<int> (n, false));
        for (int k = 0; k < 2; k ++) {
                for (int i = 0; i < n; i ++) {
                        if (g[k][i].size() & 1) zero[k][i] = true;
                }
        }
        bool ok = true;
        for (int i = 0; i < n; i ++) ok &= zero[0][i] == zero[1][i];
        if (!ok) {
                cout << "IMPOSSIBLE" << endl;
                return 0;
        }
        for (int k = 0; k < 2; k ++) dfs(root[k], k, zero);
        ans.resize(n, 0);
        for (int i = 0; i < n; i ++) {
                if (!zero[0][i] && !used[i]) { 
                        coloring(i, 1);
                }
        }
        cout << "POSSIBLE" << endl;
        for (int i = 0; i < n; i ++) cout << ans[i] << (i == n - 1 ? '\n' : ' ');
        return 0;
}