主要思路
这道题比较有趣,一开始我想出了\(O(n^3)\)的复杂度,打上去 TLE 之后思考前缀做差的方法,但是我太 Naive,不太相信(?)模意义下的减法,所以这个思路在我看了题解之后才搞定。
这道题的 DP 方程不难想到,为:
\[ dp[u][i] = \sum_{t \in dst} \sum_{j \in color(t)} dp[t][j], i \neq j \]
其实我们可以把第二维消掉,考虑总和\(tot[]\)数组,得:
\[ dp[u][i] = \sum_{t \in dst} tot[t] – dp[t][i] \]
然后就出来了。
代码
// P3914.cpp #include <bits/stdc++.h> using namespace std; const int MAX_N = 5050, MOD = 1e9 + 7; int head[MAX_N], current, n, m, tmpx, tmpy, dp[MAX_N][MAX_N], tot[MAX_N]; struct edge { int to, nxt; } edges[MAX_N << 1]; void addpath(int src, int dst) { edges[current].to = dst, edges[current].nxt = head[src]; head[src] = current++; } void dfs(int u, int fa) { for (int i = head[u]; i != -1; i = edges[i].nxt) if (edges[i].to != fa) dfs(edges[i].to, u); for (int i = 1; i <= m; i++) { if (!dp[u][i]) continue; for (int e = head[u]; e != -1; e = edges[e].nxt) { if (edges[e].to == fa) continue; dp[u][i] = (1LL * dp[u][i] * 1LL * (tot[edges[e].to] - dp[edges[e].to][i])) % MOD; } while (dp[u][i] < 0) dp[u][i] += MOD; tot[u] = (1LL * tot[u] + 1LL * dp[u][i]) % MOD; } } int main() { memset(head, -1, sizeof(head)); scanf("%d%d", &n, &m); for (int i = 1; i <= n; i++) { scanf("%d", &tmpy); for (int j = 1; j <= tmpy; j++) scanf("%d", &tmpx), dp[i][tmpx] = 1; } for (int i = 1; i < n; i++) scanf("%d%d", &tmpx, &tmpy), addpath(tmpx, tmpy), addpath(tmpy, tmpx); dfs(1, 0); printf("%d", tot[1] % MOD); return 0; }