1. 项目背景与题目解析
这道P4084 [USACO17DEC] Barn Painting G题目来自美国计算机奥林匹克竞赛(USACO)2017年12月银组真题。题目描述了一个经典的树形结构染色问题:给定一棵N个节点的树,部分节点已经被预先染色(颜色限制为红、绿、蓝三种之一),要求计算在相邻节点颜色不同的约束下,整棵树的合法染色方案数,结果对10^9+7取模。
作为信奥刷题系列的一部分,这道题很好地融合了树形DP和组合数学的思想。在实际竞赛中,这类题目常出现在省选及以上级别的比赛中,考察选手对树形结构的处理能力和动态规划的状态设计能力。我选择用C++实现不仅因为其执行效率高,更因为STL库提供的vector等容器能优雅地处理树形结构。
2. 算法核心思路
2.1 树形DP基础框架
树形动态规划的核心在于递归遍历树结构,并在回溯时完成状态转移。对于本题,我们定义dp[u][c]表示以u为根的子树,当u节点染成c颜色时的合法方案数。状态转移方程需要考虑子节点的颜色限制:
code复制dp[u][c] = Π (sum dp[v][c'] for c' ≠ c)
其中v是u的所有子节点
这个转移式的含义是:当前节点选择某种颜色后,其所有子节点选择不同颜色的方案数之积。对于预先染色的节点,只需要保留指定颜色的方案数,其他颜色方案数置零。
2.2 颜色限制处理技巧
题目允许部分节点预先染色,这需要在初始化阶段特殊处理。我的做法是:
- 读取预设颜色后,将该节点其他颜色的dp值初始化为0
- 未被预设的节点,三种颜色的dp值初始化为1
- 使用DFS后序遍历,确保子节点的状态先于父节点计算完成
2.3 模数运算注意事项
由于结果可能非常大,题目要求对1e9+7取模。这里有几个易错点:
- 乘法运算后要立即取模,防止溢出
- 建议使用long long类型存储中间结果
- 负数取模需要特殊处理(本题不涉及但值得注意)
3. C++实现详解
3.1 数据结构设计
cpp复制const int MOD = 1e9+7;
vector<vector<int>> tree; // 邻接表存储树结构
vector<array<ll,3>> dp; // dp[u][0..2] 表示三种颜色
vector<int> color; // 0表示未染色,1-3对应三种颜色
使用邻接表存储树结构比传统的指针式二叉树更灵活,能处理多叉树的情况。array容器在栈上分配,访问效率高于vector。
3.2 DFS遍历实现
cpp复制void dfs(int u, int parent) {
for(int v : tree[u]) {
if(v == parent) continue;
dfs(v, u);
for(int c = 0; c < 3; ++c) {
if(color[u] != 0 && c != color[u]-1) continue;
ll sum = 0;
for(int child_c = 0; child_c < 3; ++child_c) {
if(child_c == c) continue;
sum = (sum + dp[v][child_c]) % MOD;
}
dp[u][c] = (dp[u][c] * sum) % MOD;
}
}
}
关键点说明:
- 通过parent参数避免重复访问父节点
- 内层循环处理颜色约束条件
- 乘法累加时及时取模
3.3 主函数逻辑
cpp复制int main() {
// 输入处理
int N, K;
cin >> N >> K;
tree.resize(N+1);
dp.resize(N+1, {1,1,1});
color.resize(N+1);
// 建树
for(int i=1; i<N; ++i) {
int x, y;
cin >> x >> y;
tree[x].push_back(y);
tree[y].push_back(x);
}
// 处理预设颜色
while(K--) {
int b, c;
cin >> b >> c;
color[b] = c;
for(int i=0; i<3; ++i) {
if(i != c-1) dp[b][i] = 0;
}
}
// 执行DP
dfs(1, -1);
// 输出结果
ll ans = 0;
if(color[1] != 0) {
ans = dp[1][color[1]-1];
} else {
ans = (dp[1][0] + dp[1][1] + dp[1][2]) % MOD;
}
cout << ans << endl;
return 0;
}
4. 性能优化与边界处理
4.1 时间复杂度分析
该算法对每个节点访问一次,每个节点处理三种颜色,每种颜色需要遍历其子节点的两种颜色,因此总时间复杂度为O(N32)=O(N),完全能够处理题目中N≤1e5的数据规模。
4.2 空间优化技巧
虽然我们使用了O(N)的额外空间存储DP数组,但实际上可以进一步优化:
- 对于叶子节点,回溯后其DP值不再需要,可以释放
- 使用内存池技术复用空间
但在竞赛编程中,这种优化通常不必要,优先保证代码清晰更重要。
4.3 常见错误排查
- 忘记处理取模导致溢出:建议封装mul函数自动处理
- 建树时未处理双向边:树是无向图,邻接表需要添加双向边
- 根节点选择影响结果:本题指定1为根,无需考虑
- 预设颜色处理不完整:需要同时清除其他颜色的可能性
5. 测试用例设计
验证代码正确性的关键测试场景:
text复制// 样例输入1
4 1
1 2
1 3
1 4
2 1
// 应输出8
// 样例输入2(链式结构)
3 2
1 2
2 3
1 1
3 2
// 应输出0(无法满足约束)
// 大型随机测试
// 生成随机树和随机颜色约束,验证结果合理性
6. 同类问题扩展
掌握这道题后,可以尝试以下变种问题:
- 颜色数量扩展到K种:修改DP状态维度即可
- 相邻节点颜色差大于等于D:增加状态转移条件
- 树上最大独立集问题:将颜色视为选/不选状态
- 带权值的染色问题:在DP值中累加权值
关键心得:树形DP的难点在于设计合适的状态表示和正确的转移顺序。建议先在纸上画出小规模案例,手动模拟DP过程,再转化为代码。对于复杂的状态转移,使用注释明确每个步骤的意图非常必要。