TopCoder SRM 654 Hard

2016-03-30 | [Competitive programming]

writer 解と違う解き方っぽいですが面白いなーと思ったので書いておきます

概要

n 個の頂点からなる木構造をした家がある．n 頂点のうち 2 頂点は出入り口になっている．

n 個の家具を家に搬入し，各部屋にちょうど 1 個ずつの家具を置くようにしたい．家具は互いに区別されるが，どの家具をどの部屋に置いてもかまわない．

家具は，2 つの出入り口のどちらかから搬入し，目的の部屋までの最短路を通過して目的の部屋に置く．

家具はとても大きいので，一度家具が置かれた部屋を以降通過することはできない．

家具の置き方の場合の数を mod 1000000007 で求めよ．

解法

2 つの出入り口 s1, s2 を結ぶパスを考え，パスの上で最初に荷物が置かれる場所が p であるときの場合の数を A(p) とする．また，パスの上の 2 頂点 p, q を結ぶ辺 p–q を使用禁止にしたときの場合の数を B(p, q) とする．

すると，B(p, q) = A(p) + A(q) であることがわかる．なので，2 * (答え) = 2 * ΣA(p) (p はパス上の頂点) = A(s1) + A(s2) + ΣB(p, q) (p–q はパス上の辺) となる．

A(s1), A(s2) は，それぞれ出入り口 s1, s2 を無視して DP を行うと求められる．また，B(p, q) は s1, s2 側でほぼ独立に解けるので，これで答えの 2 倍が求まる．すると 1000000007 は奇数なので答えはただちに求まる．

コード

vector<int> graph[3030];
int N;
int C[3030][3030];
int fa, fb;
int dist[2][3030];
bool onpath[3030];

void cdist(int p, int rt, int t, int d)
{
  dist[t][p] = d;
  for (int q : graph[p]) if (q != rt) {
    cdist(q, p, t, d + 1);
  }
}

pair<i64, int> solve(int p, int rt)
{
  int size = 0;
  i64 ans = 1;
  for (int q : graph[p]) if (q != rt) {
    if (p == fa && q == fb) continue;
    if (q == fa && p == fb) continue;
    auto s = solve(q, p);
    ans = ans * s.first % MOD * C[size + s.second][size] % MOD;
    size += s.second;
  }
  return{ ans, size + 1};
}

class TwoEntrances {
  public:
  int count(vector <int> a, vector <int> b, int s1, int s2)
  {
    N = a.size() + 1;
    for (int i = 0; i <= N; ++i) C[0][i] = 0;
    C[0][0] = 1;
    for (int i = 1; i <= N; ++i) {
      C[i][0] = 1;
      for (int j = 1; j <= N; ++j) C[i][j] = (C[i - 1][j - 1] + C[i - 1][j]) % MOD;
    }

    for (int i = 0; i < N; ++i) graph[i].clear();
    for (int i = 0; i < N - 1; ++i) {
      graph[a[i]].push_back(b[i]);
      graph[b[i]].push_back(a[i]);
    }
    cdist(s1, -1, 0, 0);
    cdist(s2, -1, 1, 0);
    for (int i = 0; i < N; ++i) {
      onpath[i] = (dist[0][i] + dist[1][i] == dist[0][s2]);
    }
    i64 ret = 0;
    fa = -1; fb = -1;
    ADD(ret, solve(s1, -1).first);
    ADD(ret, solve(s2, -1).first);
    for (int p = 0; p < N; ++p) {
      for (int q : graph[p]) if (p < q && onpath[p] && onpath[q]) {
        fa = p; fb = q;
        auto a1 = solve(s1, -1), a2 = solve(s2, -1);
        ADD(ret, a1.first * a2.first % MOD * C[N][a1.second] % MOD);
      }
    }
    return (ret * ((MOD + 1) / 2)) % MOD;
  }
};