AT4289 [ABC133E] Virus Tree 2

原题链接

• 题意：给定一个 (n) 个节点的树。现在你拥有 (k) 种颜色，你要用这些颜色给树上的每个节点染色，使得任何两个距离不大于2的不同节点所被染的颜色不同。
  由于答案可能过大，请将其对 (10^9+7) 取模 $$1<=N,K <= 1e5。
• 题解：由于 (n,k) 都是 (1e5) ,并且显然的是，设 (a_i) 是表示 (i) 点能涂的颜色种类，那么最终方案数一定是 (prod_{i = 1}^n a_i)，所以就从数根向下模拟即可。
• 代码：

#include <iostream>
#include <cstring>
#include <cmath>
#include <algorithm>
#include <queue>
#include <vector>
#include <map>

#include <cstdio>

using namespace std;
typedef long long ll;
const int N = 2e5 + 99;
const ll mod =  1000000007;
vector<int>G[N];
ll a[N];
int dep[N];
struct node {
    int pos, dis;
};
queue<node>q;
ll n, k;
void dfs(int u, int fa) {
    if (fa == -1)dep[u] = 1;
    else dep[u] = dep[fa ] + 1;
    for (auto v:G[u]) {
        if (v != fa) dfs(v, u);
    }
}
void dfs1(int u, int fa) {
    if (dep[u] == 2) {
        a[u] -= 1;
    } else if (dep[u] > 2) a[u] -= 2;
    int cnt = 0;
    //if (u == 4)cout << a[u] << endl;
    for (auto v : G[u]) {
       // cout << u << ":" << v << endl;
        if (v != fa) {
            a[v] -= cnt;
            //cout << v <<"----"<< endl;
            //cout << cnt << endl;
            cnt++;
            dfs1(v, u);
            
        }
    }
}
void solve() {

    cin >> n >> k;
    for (int i = 1; i < n; i ++ ){
        int u, v;
        a[i] = k;
        cin >> u >> v;
        G[u].push_back(v);
        G[v].push_back(u);
    }
    dfs(1, -1);
    a[n] = k;
    dfs1(1, -1);
    ll ans = 1;
    for (int i = 1; i <= n; i++) 
    {
        (ans *= a[i] )%= mod;
    }
    cout << ans << endl;
}
int main()
{
    ios::sync_with_stdio(0);
    int t = 1;//cin >>t;
    while (t--)
    solve();return 0;
}