ツバサの備忘録

主に備忘録代わりに精進記録を載せていくつもりです。

SnackDown 2016 - Coloring A Tree

問題

N-1本の辺とN個の頂点からなる連結無向グラフが与えられます。それを、ある条件を満たすようにK色以下で塗り分ける場合の数を、10^{9}+7の余りで出力しなさい、というものです。
条件が、任意の頂点と同じ色の別の頂点は、ほかの同じ色の頂点のみをたどってそこにたどりつけるようにする、というものです。
これをT回繰り返します。

解法

まずは使う色の種類を固定します。i色使うとします。
このとき、木はi個に分かれていないといけません。
ここでのポイントが、辺をi-1個取り除くと、木はi個になる、ということです。
ということで、  _{N-1}C_{i-1}を計算すると、i個の分け方を求めることができます。
あとは、色の使い方を求めますが、これはK色からi色を選んで並べる順列と同じになります。
ということで、
 _{N-1}C_{i-1} × _{K}P_{i}をどんどん足していき、iを1からKまでループさせれば終わりです。
提出コードはこちらになります。

#include <bits/stdc++.h>
#define N (long long)(1e9 + 7)
#define MAX 5000
using namespace std;

long long factorial[MAX] = {0}, finverse[MAX] = {0},
          inverse[MAX] = {0};

void smodfact() {
  factorial[0] = factorial[1] = 1;
  finverse[0] = finverse[1] = 1;
  inverse[1] = 1;
  for(int i = 2; i < MAX; ++i) {
    factorial[i] = factorial[i - 1] * i % N;
    inverse[i] = N - (inverse[N % i] * (N / i)) % N;
    finverse[i] = finverse[i - 1] * inverse[i] % N;
  }
}

long long calccomb(long long n, long long k) {
  if(n < 0 || k < 0 || n < k) return 0;
  return factorial[n] * finverse[k] % N * finverse[n - k] %
         N;
}
int t, n, k, u, v;

int main() {
  long long ans = 0, now;
  smodfact();
  cin >> t;
  while(t) {
    ans = 0;
    cin >> n >> k;
    for(int i = 0; i < n - 1; ++i) cin >> u >> v;
    for(int i = 0; i < n; ++i) {
      now = calccomb(n - 1, i);
      now *= calccomb(k, i + 1);
      now %= N;
      now *= factorial[i + 1];
      now %= N;
      ans += now;
      ans %= N;
    }
    cout << ans << endl;
    --t;
  }
  return 0;
}