WARush

SRMの結果とか、解けた問題のコードを書いていきます

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);
    }
};