题目描述

难度分:$1505$

输入$n(1 \leq n \leq 10^5)$，$k(1 \leq k \leq 10^5)$和一棵无向树的$n-1$条边，节点编号从$1$到$n$。

有$k$种颜色，你需要把每个节点染成$k$种颜色中的一个。要求：对于相距为$1$或$2$的点对，颜色必须不同。

输出染色方案数，模$10^9+7$。

输入样例$1$

4 3
1 2
2 3
3 4

输出样例$1$

6

输入样例$2$

5 4
1 2
1 3
1 4
4 5

输出样例$2$

48

输入样例$3$

16 22
12 1
3 1
4 16
7 12
6 2
2 15
5 16
14 16
10 11
3 10
3 13
8 6
16 8
9 12
4 3

输出样例$3$

271414432

算法

树形DP

状态定义

$dp[u]$表示以$u$为根的子树有多少种染色方法。在这个定义下，$dp[1]$就是答案（以$1$为总根，别的节点为总根也可以）。

状态转移

从根节点开始进行DFS，逐个为所有节点染色。当前节点$u$不能和父亲节点$fa$颜色相同（如果有父亲的话），不能和祖父节点$grand\_fa$颜色相同（如果有祖父节点的话）。还不能与自己遍历过的兄弟节点颜色相同，设之前遍历过的兄弟节点有$bro$个，则$dp[u]=max(0,k-(fa \gt 0)-(grand\_fa \gt 0)-bro)$。

然后遍历$u$的子节点，$u$还需要和自己子节点$v$的子树染色方案组合，因此状态转移方程为$dp[u]=max(0,k-(fa \gt 0)-(grand\_fa \gt 0)-bro) \Pi_{v}dp[v]$。注意在遍历$v$的时候，遍历到$v$时要维护此时有多少个$v$的兄弟被遍历过了，作为$v$的$bro$。

复杂度分析

时间复杂度

本质上就是对整棵树进行一次DFS，时间复杂度为$O(n)$。

空间复杂度

树的邻接表占用空间为$O(n)$。在最差情况下整棵树退化成一根链，递归深度会达到$O(n)$级别。所以整个算法的额外空间复杂度为$O(n)$。

C++ 代码

#include <bits/stdc++.h>

using namespace std;
typedef long long LL;

const int N = 100010, MOD = 1e9 + 7;
int n, k;
vector<int> graph[N];

int dfs(int u, int fa, int grand_fa, int bro) {
    LL res = k - (fa > 0? 1: 0) - (grand_fa > 0? 1: 0) - bro;
    if(res <= 0) {
        return 0;
    }
    int bcnt = 0;
    for(int v: graph[u]) {
        if(v == fa) continue;
        res = res * dfs(v, u, fa, bcnt++) % MOD;
    }
    return res;
}

int main() {
    scanf("%d%d", &n, &k);
    for(int i = 1; i <= n; i++) {
        graph[i].clear();
    }
    for(int i = 1; i < n; i++) {
        int a, b;
        scanf("%d%d", &a, &b);
        graph[a].push_back(b);
        graph[b].push_back(a);
    }
    printf("%d\n", dfs(1, 0, 0, 0));
    return 0;
}