Virus Tree 2（树形DP）

题意

给定一个$N$个节点的树以及$K$种颜色。对于树上的每个节点，你可以选择其中一种颜色给它染色，但是必须满足如下条件：

任意两个距离不超过$2$的节点，那么的颜色不能相同。

问：有多少种染色的方案。

题目链接：https://atcoder.jp/contests/abc133/tasks/abc133_e

数据范围

$1 \leq N, K \leq 10^5$

思路

从上到下求每个节点的染色方案数，最终所有节点的方案数相乘即为答案。

根节点的染色数是$k$。

对于每个节点，对其产生影响的节点是：父亲的父亲节点、父亲节点、已经枚举过的兄弟节点。

因此，我们分类讨论，

如果是第一个枚举到的子节点，那么其染色数是$k - 1$，如果父亲的父亲节点存在，那么为$k - 2$。

如果不是第一个枚举到的子节点，那么其染色数为上一个子节点的染色数$-1$。

代码

#include <iostream>
#include <cstdio>
#include <cstring>
#include <algorithm>

using namespace std;

typedef long long ll;

const int N = 100010, M = 2 * N, mod = 1e9 + 7;

int n;
ll k;
int h[N], e[M], ne[M], idx;
ll f[N];

void add(int a, int b)
{
    e[idx] = b, ne[idx] = h[a], h[a] = idx ++;
}

void dfs(int u, int fa)
{
    ll last = 0;
    for(int i = h[u]; ~i; i = ne[i]) {
        int j = e[i];
        if(j == fa) continue;
        if(!last) {
            if(fa == -1) f[j] = k - 1;
            else f[j] = k - 2;
        }
        else f[j] = last - 1;
        last = f[j];
        dfs(j, u);
    }
}

int main()
{
    scanf("%d%lld", &n, &k);
    memset(h, -1, sizeof h);
    for(int i = 1; i < n; i ++) {
        int a, b;
        scanf("%d%d", &a, &b);
        add(a, b), add(b, a);
    }
    f[1] = k;
    dfs(1, -1);
    ll ans = 1;
    for(int i = 1; i <= n; i ++) ans = ans * f[i] % mod;
    printf("%lld\n", ans);
    return 0;
}