Virus Tree 2

atcoder.jp

DFSをしながらK色で木を塗っていく。x,yの距離が2以下ならば色が異なるということで頂点を一つ決めた時に、それの親と親に隣接するやつ以外で塗ればいいことがわかる。

#include <bits/stdc++.h>
#pragma GCC target("avx2")
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")
#define ll long long
#define int long long
#define REP(i, n) for (long long i = 0; i < (int)(n); i++)
#define rep(i, n) for (int i = 0; i < (int)(n); i++)
#define ALL(a) (a).begin(), (a).end()
#define fore(i, a) for (auto &i : a)
using namespace std;
using in = int;
using vin = vector<in>;
const ll MOD = 1e9 + 7;
in N, K;
vin E[101010];
int dfs(in cu, in pa, in rest)
{
  in res = rest;
  in my = 1;
  if (cu > 0)
    my++;
  fore(to, E[cu]) if (to != pa)
  {
    res = (res * dfs(to, cu, K - my)) % MOD;
    my++;
  }
  return res;
}
signed main()
{
  cin >> N >> K;
  rep(i, N - 1)
  {
    in a, b;
    cin >> a >> b;
    a--;
    b--;
    E[a].push_back(b);
    E[b].push_back(a);
  }
  cout << dfs(0, -1, K);
}