SRM621 Div1 Medium "TreesAnalysis"
問題
TopCoder Statistics - Problem Statement
訳
ヴァーサは2つの無向木を持っている。それぞれの木はn個の頂点を持っている。そして、それらの頂点には0からn-1の番号が、特に順序は考えられずに振られている。2つの木の形は違うこともある。
あなたは2つの木を表した、それぞれがn-1個の要素を持つint配列 tree1とtree2が与えられる。1番目の木において、番号iの頂点はtree1[i]番目の頂点と辺で繋がっている。同様に、2番目の木において、番号iの頂点はtree2[i]番目の頂点と辺で繋がっている。
ヴァーサはこれら2つの木が相似度を計算してみた。それを彼は2つのステップで行った。まず始めに、彼はそれぞれの木で1つの辺に着目し、それらを e1, e2としたときに、S(e1, e2)の値を定義した。それから、彼は(e1, e2)の全てのペアにおける、sqr(S(e1, e2))の総和を2つの木の相似度と定義した。ここでのsqr(x) は x * x を表す。
S(e1, e2)を定義するときは、下記のような手順を踏むことにした。
- 1番目の木から、辺e1を取り除くと2つの木に分割される。この中から1つ選び、それをAと呼ぶ。
- 2番目の木から、辺e2を取り除くと2つの木に分割される。この中から1つ選び、それをBと呼ぶ。
- AとBの両方に存在する頂点のセットをXとする。
そして、サイズが最も大きい頂点セットXのサイズをS(e1, e2)と定義する。(AとBの選択がXのサイズに影響する)
与えられた2つの木の相似度を返せ。
制約
2 <= それぞれの木の頂点数 <= 4000
考えたこと
今、仮にAとBを選択済みだとして、選択されていない方をA'、B'とする。
すると、計算したいのは次の4パターンとなる。
AとBを選択したときのXのサイズ AとB'を選択したときのXのサイズ A'とBを選択したときのXのサイズ A'とB'を選択したときのXのサイズ
こいつらは、どれか1つでも計算できれば簡単にだせる。
頂点数をn、Aの頂点数をa、Bの頂点数をbとする。
また、AとBを選択したときのXのサイズをxとする。
すると、4パターンそれぞれの計算方法は次のようになる。
A とB -> Xのサイズは x (これは当たり前・・) A とB' -> Xのサイズは a - x A'とB -> Xのサイズは b - x A'とB' -> Xのサイズは n - (a - x) - (b - x) - x (上の3パターンをnから引くだけ)
なので、(e1, e2)を全探索したとしても、それらでaとbとxを一瞬で出せればOKとなる。
まずはAの頂点数aを考える。
もらった1番目の木の辺情報で、どれか1つをe1として決め打ちする。それでdfsなどすれば、頂点集合Aとその頂点数a、並びにどの頂点がAに所属しているかを出せる。なのでaの取得はこれでOK。
次にBの頂点数bとxを考える。
2番目の木をdfsして、現在着目している頂点vで、vをルートとしたサブツリーをBとする。bとxはdfsでbとxを返すようにしておけば簡単に出せる。vの全子供の頂点において、dfsを呼び出し、そこで返ってきたbとxを合計しておく。後はvにおいて、bは1足せばいいし、xはvがAに含まれるなら1足せばよい。
ソースコード
int n; vector<int> T1[4040], T2[4040]; bool used[4040]; int a; long long result; class TreesAnalysis { public: long long treeSimilarity(vector <int> tree1, vector <int> tree2) { n = tree1.size() + 1; // グラフ作成 for (int i = 0; i < n - 1; i++) { T1[tree1[i]].push_back(i); T1[i].push_back(tree1[i]); T2[tree2[i]].push_back(i); T2[i].push_back(tree2[i]); } result = 0; // tree1の頂点で回す for (int i = 0; i < n - 1; i++) { // tree1を2つに分ける memset(used, false, sizeof(used)); a = 0; dfs1(0, -1, i, tree1[i]); // tree2を2つに分け、カウント dfs2(0, -1); } return result; } void dfs1(int v, int p, int x, int y) { used[v] = true; a++; for (int i = 0; i < T1[v].size(); i++) { int c = T1[v][i]; if (c == p) continue; if ((v == x && c == y) || (v == y && c == x)) continue; dfs1(c, v, x, y); } } pair<long long, long long> dfs2(int v, int p) { long long b = 0; // Bのサイズ long long x = 0; // A * BにおけるXのサイズ // 子供をまずやる for (int i = 0; i < T2[v].size(); i++) { if (T2[v][i] == p) continue; pair<int, int> r = dfs2(T2[v][i], v); b += r.first; x += r.second; } // 自分の分 b++; if (used[v]) x++; // A * B long long res = x; // A * B' res = max(res, a - x); // A' * B res = max(res, b - x); // A' * B' res = max(res, n - x - (a - x) - (b - x)); if (p != -1) result += res * res; return make_pair(b, x); } };