TCO15 Round 1B hard:TheTreeAndMan

方針は結構すぐ立ったけど実装にものすごく時間かかりました。

問題:TopCoder Statistics - Problem Statement

解法:写真を参照。

vとuが写真のような関係になっていれば(図にかかれている頂点は隣接している必要はない),
(A)vの子でuに繋がっていないもののうち,条件を満たすように頂点を2個選ぶ場合の数(写真の図で言うと,vの左下に頂点が2個あったとしてもvの左下から頂点を2個選んだものは条件に合っていないので数えてはいけない)
(B)uの子のうち条件を満たすように頂点を2個選ぶ場合の数(上と同様に左下/右下といった同じグループから頂点を2つ選んではいけない)
(C)vの上にある頂点の数
がわかれば,これらの積をすべての組み合わせ(u,v)について加えたものが答えになる。

これを実行するために,以下の作業をする。
・各頂点vについてその頂点の下にある頂点の数size[v]を数える
・各頂点vから頂点uに行くことができるかどうかのフラグcan[v][u]をたてる
・上のcanの情報を使って各頂点vに向かうことのできる頂点数rsize[v]を計算する

これがわかると,まず上の(C)はrsize[v]によってわかる。
また,(B)については例えばuの子のそれぞれのサイズが
a0, a1, a2, ..., an
となっていたとすると,次のようにして条件に合う頂点の選び方を数えられる。
・a0から1個選んだとすると残りの選び方はa1+a2+a3+...+an
・a1から1個選んだとすると残りの選び方は a2+a3+...+an
・a2から1個選んだとすると残りの選び方は a3+...+an

もちろんこのまま計算すると計算量が多くなるのでsum[]みたいな配列を作って工夫する。(A)も同様。
最悪O(n^3)なんじゃないかと思ってちょっとびびってたんですが最悪でもおよそ100msだったので大丈夫そうですね。
以下ソースコード

const int MAXN = 2222;
int size[MAXN], rsize[MAXN];
vector<int> T[MAXN];
bool can[MAXN][MAXN];
const ll MOD = 1e9+7;

int dfs(int v) {
    if (size[v] >= 0) return size[v];
    size[v] = 1;
    for (int u : T[v]) {
        size[v] += dfs(u);
    }
    return size[v];
}

void dfs2(int s, int v) {
    can[s][v] = true;
    for (int u : T[v]) {
        dfs2(s, u);
    }
}

ll totalSize(int v, int u) {
    vector<int> tmp;
    for (int el : T[v]) {
        if (u >= 0 && can[el][u]) continue;
        tmp.push_back(size[el]);
    }
    int n = tmp.size();
    vector<int> sum(n+1);
    for (int i = 0; i < n; i++) sum[i+1] += sum[i]+tmp[i];
    ll plus = 0;
    for (int i = 0; i < n; i++) {
        plus += tmp[i] * (sum[n]-sum[i+1]);
    }
    return plus;
}

class TheTreeAndMan {
public:
    int solve(vector <int> parent) {
        for (int i = 0; i < MAXN; i++) T[i].clear();
        int N = parent.size()+1;
        for (int i = 1; i < N; i++) {
            T[parent[i-1]].push_back(i);
        }

        memset(size, -1, sizeof(size));
        dfs(0);

        memset(can, 0, sizeof(can));
        for (int i = 0; i < N; i++) {
            dfs2(i, i);
        }

        for (int i = 0; i < N; i++) {
            rsize[i] = 0;
            for (int j = 0; j < N; j++) {
                if (can[j][i] && i != j) rsize[i]++;
            }
        }

        ll ans = 0;
        cout << "TEST" << endl;
        for (int v = 1; v < N; v++) {
            for (int u = 1; u < N; u++) {
                if (!can[v][u] || v == u) continue;
                ll tmp = totalSize(v, u), tmp2 = totalSize(u, -1);
                ans += (tmp * tmp2) % MOD * rsize[v];
            }
            ans %= MOD;
        }
        return ans;
    }
};