打卡刷题进入第2783天恰好撞上编号也是2783的题目算是刷题生涯里的一点小默契。P3914《染色计数》是一道经典的树形DP入门题题面很短一句话就能讲完但真正动手写的时候卡住我的不是不会树形DP而是没想清楚一个关键问题每个节点可用的颜色集合不一样怎么把状态和可选颜色干净地对齐。这篇就把我从读题、建模、优化到实现的完整思路整理出来给一起刷信奥题的朋友做个参考也适合刚学树形DP、准备从模板题往综合题过渡的读者。1. 先看懂P3914在问什么一个节点一个颜色集合的树染色计数1.1 题面拆成三句话这道题的核心信息可以浓缩成三个部分。第一给了一棵n个节点的树。树的性质很重要没有环任意两个点之间只有一条简单路径这意味着我们可以用递归的方式自底向上处理。第二有m种颜色但每个节点并不是所有颜色都能用。题目会给每个点一个可选颜色集合比如节点1只能用颜色2和3节点2只能用颜色1这就让问题比每个点都能染任意颜色复杂了一个档次。很多树形DP模板题默认颜色全集对每个点开放这道题偏偏在每个点上做了限制处理方式必须跟着变。第三相邻两个节点不能染同一种颜色问整棵树有多少种合法的染色方案答案对某个大质数取模。题面看起来人畜无害但方案计数这四个字意味着答案可能非常庞大不能直接枚举必须用动态规划把子问题的结果复用到父问题上。1.2 为什么不能暴力枚举假设每个节点平均有k种可选颜色那么暴力枚举所有染色方案的复杂度是k的n次方级别。哪怕k只有2、n只有30方案数也已经超过10亿更别说n稍微大一点或者某些点颜色集合比较大。所以要换个思路树上只有相邻关系产生约束而且树的结构决定了每个子树之间是隔着当前节点互不干扰的。把当前节点的颜色固定以后它的每个孩子子树内部怎么染色彼此之间完全独立只需要回避孩子不能和当前节点同色这一条规则。这种固定一点子树自治的特征刚好是树形DP能下手的地方。1.3 难度定位与前置知识这道题在网络上的难度标签大约是普及/提高但它考的东西比普通模板题要全状态设计、乘法原理、针对颜色集合的稀疏存储、取模运算再加上一点用二分查找优化索引的功夫。前置知识需要会建树、会DFS递归、理解树上组合计数里加法和乘法的区别。如果你已经能独立写出没有上司的舞会这类最基础的树形DP那这道题就是你从入门跨向进阶的一道好关卡。2. 状态与转移dp[u][c]是怎么推出来的为什么要按颜色拆开2.1 只存dp[u]够不够第一次做树上染色计数的人很容易下意识写一个dp[u]表示以u为根的子树染色方案数。但这个状态在转移时站不住脚当我们要合并孩子v的方案时v能不能选某个颜色取决于u最后选了哪个颜色。如果dp[u]只是一个汇总后的数字孩子的方案没法针对u具体染了什么颜色做排除。所以状态必须带上颜色维度。定义dp[u][c]为以u为根的子树内部在u被染成颜色c的前提下子树里所有节点都满足相邻不同色的合法方案数。注意这里有个隐含条件c必须属于u自己的可选颜色集合否则这个状态没有意义直接视为0即可。每个节点维护一串和可选颜色一一对应的dp值这就是后面所有优化的基础。理解这一点再看网上那些二维数组开到底的写法就明白它们其实是在给每个点都开一个长度为m的数组即使这个点根本不能用某种颜色数组位置也空着。2.2 叶子节点和内部节点的差异处理叶子节点最简单如果u是叶子它没有孩子那么当u染成c时子树内没有其他节点需要考虑约束方案数就是1。所以只要c属于u的可选颜色dp[u][c] 1。内部节点就要做乘法了。假设u染成了颜色c我们来看任意一个孩子v。v不能染c但可以染它可选颜色集合里的其他任何颜色x。当v染x时以v为根的子树的合法方案数是dp[v][x]。所以在u染c的前提下孩子v这一支能提供的方案数就是所有x不等于c时的dp[v][x]之和。不同孩子之间是并列关系各自独立所以要把每个孩子的贡献值乘起来得到dp[u][c]。2.3 加法和乘法到底怎么区分学树形DP很容易把加法和乘法搞混这里我多说几句。在同一个节点上选择不同的颜色是互斥的或关系u要么染红要么染蓝不可能同时染两个。所以最后汇总一个节点的答案时要把该节点所有颜色的dp值加起来。在不同孩子之间是并且的关系孩子甲的方案确定后孩子乙的方案可以独立选择两者组合起来才是一棵完整子树的一种方案。甲有2种乙有3种组合起来就是2乘3等于6种。把这两条规则写进代码里就是计算dp[u][c]时用乘法连接各个孩子最后计算sum[u]时用加法汇总所有颜色。2.4 用一条链手推一遍理论说多了容易飘拿个具体例子走一遍。考虑一条三个点的链1-2-3每个点可选颜色都是{1,2}。从叶子3开始dp[3][1]1dp[3][2]1sum[3]2。节点2有两个可选颜色。当2染1时孩子3不能染1只能染2贡献是dp[3][2]1所以dp[2][1]1。同理dp[2][2]1sum[2]2。节点1接两个孩子不这里节点1只有孩子2。当1染1时2不能染1贡献是sum[2]-dp[2][1]2-11所以dp[1][1]1dp[1][2]1sum[1]2。这棵树总共2种染色方案分别是(1,2,1)和(2,1,2)和手算结果一致。虽然这个例子太简单但它把排除孩子同色的过程完整展示出来了。2.5 状态空间怎么存dp[u][c]里面的c只会在u的可选颜色集合里取。所以完全不需要开一个n乘m的二维数组只需要让dp[u]是一个和col[u]等长的数组下标一一对应。这样内存是每个点实际可用颜色数量之和比n乘m小得多。真正写代码时这种稀疏存储是必须的——因为你不知道出题人会给你多大的m一旦m到1e5量级二维数组直接爆内存。3. 优化点把枚举另一方颜色换成总和减单色一次减法干掉O(m)3.1 朴素写法的复杂度隐患如果顺着转移公式直接写计算dp[u][c]时对每个孩子v都要枚举v的所有可选颜色x把它们加起来再排除x等于c的那一项。伪代码大概是for (int c : col[u]) { dp[u][c] 1; for (int v : children[u]) { int total 0; for (int x : col[v]) { if (x ! c) total dp[v][x]; } dp[u][c] dp[u][c] * total % MOD; } }这个三层循环的复杂度是每个点的颜色数乘孩子数乘孩子颜色数总计可以到O(n乘以m平方)的级别。当m到几百几千n再大一点时间完全撑不住。我最初写这道题时就是先写了这个朴素版本结果一提交就发现跑不动。问题不在于DP思路错而在于排除一种颜色这个动作做得太笨了——明明孩子的所有颜色只需要求和一次我却对每个c都重新扫一遍孩子的颜色。3.2 核心优化先求总和再做减法换个角度看孩子v的所有可选颜色的dp值之和其实是一个固定数字跟u选什么颜色没关系。所以可以提前算好sum[v] Σ dp[v][x]。这样一来u染c时孩子v的贡献就是sum[v]减去dp[v][c]如果c属于v的可选集合否则就是完整的sum[v]。注意第二句话如果c根本不在v的可选颜色集合里说明v本来就不可能染c那么u就算染c也不会对v产生任何约束v这一支的所有方案都可以保留贡献就是sum[v]。从公式上看dp[v][c]不存在视作0sum[v]-0sum[v]完全一致。这个优化把遍历孩子所有颜色做排除降成了判断一个颜色是否存在于孩子集合里。外层对col[u]的枚举仍然保留但内层不再嵌套枚举col[v]。于是总复杂度降到大约O(所有点可选颜色数总和)每条边对应一次对col[u]的遍历和m的大小没有直接关系了。3.3 一段很容易误解的公式网上很多题解会把转移写成dp[u][c] Π(sum[v] - dp[v][c])这个写法省略了一次判断看起来很简洁但对新手不太友好如果c不在v的可选集合里dp[v][c]是什么严谨的做法是先判断c在不在v的可选集合里。在就用sum[v]-dp[v][c]不在就用sum[v]。代码实现里我选择用二分查找来判断虽然多了一个log的常数但逻辑非常清晰不容易错。3.4 取模运算里的负数sum[v]和dp[v][c]都是取模后的数sum[v]减去dp[v][c]有可能得到负数。在C里负数的取模运算结果仍然是负数直接乘进答案里就错了。所以减法之后要立刻调整成非负ll factor (sum[v] - dp[v][j] MOD) % MOD;这里用long long承接乘法避免两个int相乘溢出。很多第一次写模运算DP的人会在这两个地方各交一次学费后面第五章我再集中盘点。4. 完整C实现压缩状态的树形DP代码与注释4.1 数据结构怎么选我的做法是三个核心容器col[u]存节点u的可选颜色读入后排序、去重dp[u]和col[u]等长的数组dp[u][i]对应颜色col[u][i]的方案数sum[u]dp[u]所有元素求和取模。为什么排序去重因为如果有重复的颜色出现在同一行里它们不会产生新的染色方案留着只会让dp数组和颜色索引对不上二分查找也会出问题。排序则让判断某个颜色是否在某节点的集合里可以用lower_bound做二分。为什么不用二维数组dp[n1][m1]当m比较小比如m小于等于2000二维数组确实直观查颜色直接按下标走。但一旦m到1e5空间就是n乘1e5根本开不出来。压缩存储以后内存只和每个点实际可用的颜色总数相关什么数据范围都扛得住。4.2 完整代码#include bits/stdc.h using namespace std; typedef long long ll; const int MOD 1000000007; int n, m; vectorvectorint g; // 邻接表 vectorvectorint col; // col[u] 排序去重后的可选颜色 vectorvectorll dp; // dp[u][i] 与 col[u][i] 一一对应 vectorll sum; // sum[u] Σ dp[u][i] % MOD // 二分查找 x 在数组 a 中的下标不存在返回 -1 int idIn(const vectorint a, int x) { auto it lower_bound(a.begin(), a.end(), x); if (it ! a.end() *it x) return (int)(it - a.begin()); return -1; } void dfs(int u, int fa) { int k (int)col[u].size(); dp[u].assign(k, 0); // 当前节点没有任何可选颜色整棵子树无解 if (k 0) { sum[u] 0; return; } // 初值 1表示每个孩子还没乘进来时的空乘积 dp[u].assign(k, 1); for (int v : g[u]) { if (v fa) continue; dfs(v, u); // 孩子子树已经无解当前节点所有颜色方案都为 0 if (sum[v] 0) { dp[u].assign(k, 0); sum[u] 0; return; } for (int i 0; i k; i) { int c col[u][i]; int j idIn(col[v], c); ll factor; if (j -1) { // 孩子根本不能用颜色 c不受约束 factor sum[v]; } else { // 孩子不能用 c所以从总和中减去孩子染 c 的方案数 factor (sum[v] - dp[v][j] MOD) % MOD; } dp[u][i] dp[u][i] * factor % MOD; } } sum[u] 0; for (int i 0; i k; i) { sum[u] (sum[u] dp[u][i]) % MOD; } } int main() { ios::sync_with_stdio(false); cin.tie(nullptr); cin n m; g.assign(n 1, {}); col.assign(n 1, {}); for (int i 1; i n; i) { int cnt; cin cnt; col[i].resize(cnt); for (int j 0; j cnt; j) { cin col[i][j]; } sort(col[i].begin(), col[i].end()); col[i].erase(unique(col[i].begin(), col[i].end()), col[i].end()); } // 读入 n-1 条边建无向图 for (int i 1; i n; i) { int u, v; cin u v; g[u].push_back(v); g[v].push_back(u); } dp.assign(n 1, {}); sum.assign(n 1, 0); // 树是无根树选任意节点当根都可以这里选 1 dfs(1, 0); cout sum[1] % MOD \n; return 0; }4.3 代码里几个容易看漏的细节第一读取颜色时每行第一个数cnt表示这一行后面跟着多少个颜色不是行号。这个容易和后面的n-1行边搞混我有一段时间每次读入都把节点编号当成颜色个数读错了半天才反应过来。第二把根固定为1不影响答案。因为染色约束只跟边有关和哪个点是根没有关系。任意指定根以后父子关系只是决定了DFS的调用顺序不会改变相邻约束。第三孩子sum[v]为0时当前节点所有dp直接清0并返回。这是剪枝理论上就算不剪枝后面所有乘法因子也都会变成0。加上这个判断让代码在无解数据上跑得更快也让逻辑更清晰。第四如果题面给的输入不是无向边而是每个点的父节点编号处理方式更简单把父边当成有向边建图即可DFS时也就不用判断fa了。核心公式完全不变。4.4 和二维数组写法对比有些题解开long long dp[1005][1005]然后直接dp[u][color]访问这种写法在m很小时非常痛快。核心转移只有几行for (int i 0; i k; i) { int c col[u][i]; dp[u][c] 1; for (int v : child[u]) { long long sum 0; for (int x : col[v]) if (x ! c) sum dp[v][x]; dp[u][c] dp[u][c] * sum % MOD; } }这个版本没有做总和减单色的优化其实还是朴素写法。要是直接在此基础上加sum数组优化二维数组版本也可以很快。但它的问题始终是空间m稍微一大n乘m的二维long long数组就会击穿内存。我建议平时练题就习惯压缩存储的写法将来遇到m为1e5甚至更大的题不会措手不及。5. 边界用例与实战复盘样例手算、空集合、模运算和爆栈5.1 用一个三节点样例完整手算写代码不能光靠感觉要对着一组小数据把dp表一步步推出来。下面这个例子是我当时用来验证代码正确性的节点数n3颜色数m3。 节点1可选颜色{1,2}节点2可选颜色{1,3}节点3可选颜色{2,3}。 边1-21-3。先算叶子。节点2没有孩子dp[2][1]1dp[2][3]1sum[2]2。节点3同理dp[3][2]1dp[3][3]1sum[3]2。再看节点1。对颜色1来说孩子2可选集合里有1因子是sum[2]-dp[2][1]2-11孩子3可选集合里没有1因子是sum[3]2。所以dp[1][1]1乘1乘22。对颜色2来说孩子2集合里没有2因子sum[2]2孩子3集合里有2因子sum[3]-dp[3][2]1。所以dp[1][2]1乘2乘12。最后sum[1]dp[1][1]dp[1][2]4。手动枚举验证一下1染1时2只能染33可以染2或3共2种1染2时2可以染1或33只能染3共2种。总数4完全一致。跑这段代码输出也是4说明核心转移没问题。5.2 空可选集和必然无解如果某个节点一行输入只有0也就是它没有任何可选颜色那么整个树一定无解最终输出0。代码里处理这个情况的路径有两条一是当前节点k0时直接把sum[u]设成0返回二是某个孩子无条件后父节点把所有dp置0。这两条路最终会把0一路传到根。要注意的是一旦发现孩子sum0就置当前节点为0并return这个剪枝不会影响正确性因为0乘任何数都是0等于提前结束了无效计算。5.3 单节点和全同色链单节点树是最小边界。比如n1m5节点1可选颜色{2,4}。那么方案数就是2因为只有一个点染2或染4都合法。代码里叶子节点的sum就是去重后的颜色个数输出正确。另一个我常拿来测的边界是一条链所有节点只有颜色1可选m1。如果链只有一个点答案是1只要链长大于1相邻两个点都会被迫同色答案直接变成0。代码跑起来的过程很典型叶子sum1父节点算孩子因子时sum[v]-dp[v][1]1-10于是dp[父][1]0sum[父]0逐层向上传播最终根输出0。5.4 实战中反复踩过的几个坑模运算负数是最常见的一个。sum[v]-dp[v][j]在数学上可能是负数C取模后也还是负数直接参与乘法会得到一个看似随机的错误答案。我的习惯是凡涉及减法取模立刻写成(x % MOD MOD) % MOD不要等到出错再回头改。第二个坑是乘法溢出。两个int相乘已经可能超过int范围取模前先转成long long。我在代码里把dp和sum都声明成了vectorlong long乘法表达式里的因子自动提升为long long这样就安全了。第三个坑是递归深度。当树的形态是一条长链n到几万甚至更大时递归DFS有爆栈风险。如果遇到这种情况可以把递归改成显式栈做后序遍历先处理所有孩子再处理当前节点转移公式完全不变。做OJ题时我会先看一眼题目的n上限再决定要不要用迭代版。第四个坑是读入问题。题目给的n行颜色集合每行以颜色个数开头这个格式和普通第一行n m后面n行每行若干个数字是有区别的。我用cin cnt先把个数读掉再循环cnt次读取具体颜色顺序上不要颠倒。5.5 做完这道题的一点体会对我来说P3914最大的价值不是学会了树形DP而是把三个通用技巧练扎实了判断状态需不需要加维度、用整体总和减单项来避免二层枚举、以及用压缩存储应对不确定的数据范围。尤其是第三点很多树形DP题目在状态设计上都喜欢搞每个点只有少量可选值这种设定压缩存储配合二分查找几乎是标准解法。如果你正在刷题打卡建议不要只盯着AC试着把这份代码的时间复杂度算清楚再想想如果题目变成每个点不能和相邻点以及次相邻点同色状态要怎么扩展。把一道题吃透比泛泛刷十道题有用得多。
