読者です 読者をやめる 読者になる 読者になる

Typical DP Contest N: 木

しばらく解けなくて悩んでいた問題。
N: 木 - Typical DP Contest | AtCoder

 塗られた辺の連結を保ちながら木を塗る方法は何通りかという超シンプルな問題。
 とりあえず、先に塗る辺aを固定してしまいましょう。
 それで辺aの右側がx通り、y通りだとしたら?
 ……
 と、これだけでは組み合わせは計算できないので、辺の本数も知りたいです。
 辺aの右側がx通りで、左がy通り、辺の本数は右側がsx通りで左がsyとおりだとします。
 このとき、x, yの各組み合わせに辺の並べ方(sx+syCsx通り)をかけると場合の数が計算できます。
 部分木についても同じように計算できるので、これを部分木を辿って計算していけばいいです。子がいくつもあるときは、分岐1と2をあわせて一つの分岐だと考え、…というのを繰り返せば大丈夫なので結局特に問題はないという。

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

using i64 = long long;
const i64 MOD = 1e9 + 7;

i64 comb[1010][1010];
void setComb(const int N = 1010) {
    comb[0][0] = 1;
    for(int i : _in(1, N)) {
        comb[i][0] = 1;
        for(int j : _in(1, i + 1))
          comb[i][j] = (comb[i - 1][j - 1] + comb[i - 1][j]) % MOD;
    }
}

typedef pair<i64, int> part;
part merge(const part& a, const part& b) {
    i64 comA, comB; int szA, szB;
    tie(comA, szA) = a;
    tie(comB, szB) = b;
    i64 prod = comb[szA + szB][szA] * comB % MOD;
    return part(comA * prod % MOD, szA + szB);
}

vector<vector<int>> tree;
part memo[1010][1010];
bool used[1010][1010];

part dfs(int cur_v, int par) {
    part& res = memo[cur_v][par];
    if(used[cur_v][par]) return res;
    used[cur_v][par] = true;
    res = part(1LL, 0);
    for(int nv : tree[cur_v]) if(nv != par) {
        auto temp = dfs(nv, cur_v);
        ++temp.second;
        res = merge(res, temp);
    }
    return res;
}

i64 calc(int u, int v) {
    part a = dfs(u, v);
    part b = dfs(v, u);
    a = merge(a, b);
    return a.first;
}
int main() {
    setComb();
    cin.tie(0);
    ios::sync_with_stdio(false);
    int N;
    while(cin >> N && N) {
        tree.assign(N, vector<int>());
        vector<pair<int, int>> edges(N - 1);
        for(int i : _in(N - 1)) {
            int a, b;
            cin >> a >> b;
            --a; --b;
            tree[a].emplace_back(b);
            tree[b].emplace_back(a);
            edges[i] = make_pair(a, b);
        }
        memset(used, 0, sizeof(used));
        i64 ans = 0LL;
        for(const auto& e : edges)
           ans = (ans + calc(e.first, e.second)) % MOD;
        cout << ans << endl;
    }
    return 0;
}

 ちなみに、面倒なことをせずにiを頂点とする根付き木としてdfsするのを繰り返した場合は、ある辺に対しその両端の頂点で二回、その辺を最初に塗る場合を数えているので最後に2で割ればOK。あと頂点をみているDFSと辺を見ているDFSではサイズの渡し方が異なることに注意。

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

using i64 = long long;
const i64 MOD = 1e9 + 7;

template<typename T>
T mod_pow(T x, T n, const T& mod) {
    T res = 1;
    for(; n > 0; n >>= 1) {
        if(n & 1) res = res * x % mod;
        x = x * x % mod;
    }
    return res;
}

template <typename T>
T mod_inv(T x, T mod) {
    return mod_pow(x, mod - 2, mod);
}

i64 comb[1010][1010];
void setComb(const int N = 1010) {
    comb[0][0] = 1;
    for(int i : _in(1, N)) {
        comb[i][0] = 1;
        for(int j : _in(1, i + 1))
          comb[i][j] = (comb[i - 1][j - 1] + comb[i - 1][j]) % MOD;
    }
}

typedef pair<i64, int> part;
part merge(const part& a, const part& b) {
    i64 comA, comB; int szA, szB;
    tie(comA, szA) = a;
    tie(comB, szB) = b;
    i64 prod = comb[szA + szB][szA] * comB % MOD;
    return part(comA * prod % MOD, szA + szB);
}

vector<vector<int>> tree;
part memo[1010][1010];
bool used[1010][1010];

part dfs(int cur_v, int par) {
    /*
    part& res = memo[cur_v][par];
    if(used[cur_v][par]) return res;
    used[cur_v][par] = true;
    res = part(1LL, 0);
      */
    part res(1LL, 0);
    for(int nv : tree[cur_v])
      if(nv != par) 
        res = merge(res, dfs(nv, cur_v));
    res.second++;
    return res;
}

int main() {
    setComb();
    cin.tie(0);
    ios::sync_with_stdio(false);
    int N;
    while(cin >> N && N) {
        tree.assign(N, vector<int>());
        for(int _ : _in(N - 1)) {
            int a, b;
            cin >> a >> b;
            --a; --b;
            tree[a].emplace_back(b);
            tree[b].emplace_back(a);
        }
//        memset(used, 0, sizeof(used));
        i64 ans = 0LL;
        for(int v : _in(N))
           ans = (ans + dfs(v, -1).first) % MOD;
        ans = (ans * mod_inv(2LL, MOD)) % MOD;
        cout << ans << endl;
    }
    return 0;
}

↑でコメントアウトしてあるのはメモ化の部分ですがこれを削るとかえって速くなりました。
なんでだろう……。