とても教育的なツリーの探索問題。

問題概要

$N$ 頂点の重み付き無向グラフが与えられる。このグラフの各頂点を白黒に塗る方法であって

同色に塗られた 2 頂点はどれを選んでも、その間のパスの長さは偶数である

という条件を満たすものを求めよ。

制約

$1 \le N \le 10^{5}$

考えたこと

考え方としては

二部グラフ判定

に近い。今回は各辺に重みがあるが、この重さは偶数か奇数かだけが問題である。

その重さが奇数だったら、長さが 1 の辺
その重さが偶数だったら、長さが 0 の辺 (あるいは u-v だとすると、u-w-v と間に頂点 w が挿入された長さ 2 の状態)

という感じ。今回はツリーなので「ツリー上の DFS」というのと「二部グラフ判定」というのがいい感じに組み合わせるような実装になる。

ツリーを探索するとは

ツリーというと根から伸びているものを思いがちだが、今回は特に根が指定されているわけではない。そこでツリーを探索するときの常套手段は、

根をなんでもいいから 1 個決めてしまう

というもの。根を決めてしまえば、そこから自動的に探索することができる。特に理由がなければ根はノード 0 で構わない。そして多くの場合、以下のような書き方をテンプレとして持つことができる。

実際のツリー上の探索はこれにいろんなものを付け加えていく感じだが、ほとんどのツリー探索は下のような実装にちょっと付け加えるだけで書くことができる。

下の実装で「if (nv == p) continue;」のところがツリー探索で一般的な枝刈りで、これを行うことで、根から葉に向かって逆流することなく探索を行うことができる。

using Graph = vector<vector<int> >;
Graph G;

// v は現在見ている頂点, p は v の親
void dfs(int v, int p ) {
  for (auto nv : G[v]) {
    if (nv == p) continue;  // これがツリー探索で一般的な書き方
    dfs(nv, v);
  }
}

int main() {
  int root = 0;
  dfs(root, -1); // root の親はいないので -1
}

今回のツリー探索関数

今回のツリー探索は「色を塗る」という目的があるので、それに沿った設計にする。

下のように、再帰関数の引数に色を表す変数 c を付け加えてあげて、「次のノードを再帰的に呼び出す」というのをするときに

エッジの重みが偶数だったら、次のノードも同色 c で塗る
エッジの重みが奇数だったら、次のノードは異色 1-c で塗る

という風にすれば OK。

// 入力
Graph G;

// res は答え
vector<int> res;

// c: 0 か 1 か。v を c に塗る
void dfs(int v, int p, int c) {
    res[v] = c;
    for (auto e : G[v]) {
        if (e.first == p) continue;
        if (e.second & 1) dfs(e.first, v, 1-c);
        else dfs(e.first, v, c);
    }
}

コード

以上をまとめて $O(N)$ の計算時間でできる。

#include <iostream>
#include <vector>
using namespace std;

using Edge = pair<int,int>;
using Graph = vector<vector<Edge> >;

// 入力
int N;
Graph G;

// res は答え
vector<int> res;

// v を c に塗る。p は v の親
void dfs(int v, int p, int c) {
    res[v] = c;
    for (auto e : G[v]) {
        if (e.first == p) continue;
        if (e.second & 1) dfs(e.first, v, 1-c);
        else dfs(e.first, v, c);
    }
}

int main() {
    cin >> N;
    G.assign(N, vector<Edge>());
    for (int i = 0; i < N-1; ++i) {
        int u, v, w; cin >> u >> v >> w;
        --u, --v;
        G[u].push_back(Edge(v, w));
        G[v].push_back(Edge(u, w));
    }
    res.assign(N, 0);
    dfs(0, -1, 1);
    for (auto v : res) cout << v << endl;
}