【题解】 「NOI2020」命运 树形dp+线段树合并 LOJ3340
Legend
给定一棵 \(n\) 节点的树，你需要把边黑白染色。当然还有 \(m\) 个限制，限制是一条简单路径，且两端满足祖先-后代关系，表示这条路径上所有的边至少有一条是黑色的，问方案数对 \(998 244 353\) 取模的值。
\(1 \le n \le 500000\)，\(0 \le m \le 500000\)。
Editorial
两端满足祖先-后代关系这个条件很奇怪，不妨从此下手。
设 \(dp_{i,j}\) 表示从子树 \(i\) 里向上伸出来的链 **上端点深度最大是 \(j\) **的方案数量。
现在考虑到了节点 \(i\)，看看加入一棵子树 \(k\) 发生了什么：
\(dp'_{i,j}=\sum\limits_{s=0}^{j} dp_{i,j}\times dp_{k,s} + \sum\limits_{s=0}^{j-1} dp_{i,s} \times dp_{k,j}+ \sum\limits_{s=0}^{dep_{i}}dp_{i,j} \times dp_{k,s}\)。
其中前两个转移是 \((i,k)\) 这条边没选的，最后一个是这条边选了的。
这样就可以 \(O(n)\) 进行转移。可以写出一个 \(O(n \cdot \rm{maxdep})\) 的做法。
把方程写成前缀和形式：
\[\begin{aligned}
dp'_{i,j}&=dp_{i,j}\times S_{k,j} + dp_{k,j} \times S_{i,j-1} + dp_{i,j} \times S_{k,dep_i} \\
&=dp_{i,j}\times(S_{k,j}+S_{k,dep_i})+dp_{k,j}\times S_{i,j-1}
\end{aligned}
\]
考虑怎么快速算这一坨东西。发现前缀和其实可以通过线段树合并的时候顺带求出来。
于是复杂度就变成了 \(O(n \log n)\)。
Code
算是个套路，但是第一次见还是很有趣的。
#include <bits/stdc++.h>
#define LL long long
#define debug(...) fprintf(stderr ,__VA_ARGS__)
#define __FILE(x)\
freopen(#x".in" ,"r" ,stdin);\
freopen(#x".out" ,"w" ,stdout
const int MX = 5e5 + 233;
const LL MOD = 998244353;
int read(){
char k = getchar(); int x = 0;
while(k < '0' || k > '9') k = getchar();
while(k >= '0' && k <= '9') x = x * 10 + k - '0' ,k = getchar();
return x;
}
std::vector<int> limit[MX];
int head[MX] ,tot ,n;
struct edge{
int node ,next;
}h[MX << 1];
void addedge(int u ,int v ,int flg = 1){
h[++tot] = (edge){v ,head[u]} ,head[u] = tot;
if(flg) addedge(v ,u ,0);
}
struct node{
int l ,r;
LL sum ,mul;
node *lch ,*rch;
}*root[MX];
node *newnode(int l ,int r){
node *x = new node;
x->l = l ,x->r = r;
x->sum = 0 ,x->mul = 1;
x->lch = x->rch = nullptr;
return x;
}
void domul(node *x ,LL v){
x->sum = x->sum * v % MOD;
x->mul = x->mul * v % MOD;
}
void pushdown(node *x){
if(x->mul != 1){
if(x->lch != nullptr) domul(x->lch ,x->mul);
if(x->rch != nullptr) domul(x->rch ,x->mul);
x->mul = 1;
}
}
void pushup(node *x){
x->sum = 0;
if(x->lch != nullptr) x->sum = x->lch->sum;
if(x->rch != nullptr) x->sum = (x->sum + x->rch->sum) % MOD;
}
void change(node *x ,int p ,int val){
if(x->l == x->r) return x->sum = val ,void();
int mid = (x->l + x->r) >> 1;
pushdown(x);
if(p <= mid){
if(x->lch == nullptr) x->lch = newnode(x->l ,mid);
change(x->lch ,p ,val);
}else{
if(x->rch == nullptr) x->rch = newnode(mid + 1 ,x->r);
change(x->rch ,p ,val);
}return pushup(x);
}
LL sum(node *x ,int l ,int r){
if(x == nullptr) return 0;
if(l <= x->l && x->r <= r) return x->sum;
pushdown(x);
int mid = (x->l + x->r) >> 1;
LL s = 0;
if(l <= mid) s = sum(x->lch ,l ,r);
if(r > mid) s = (s + sum(x->rch ,l ,r)) % MOD;
return s;
}
node *merge(node *x ,node *y ,LL &s1 ,LL &s2){
if(x == nullptr){
if(y != nullptr){
s1 = (s1 + y->sum) % MOD;
domul(y ,s2);
}
return y;
}
if(y == nullptr){
s2 = (s2 + x->sum) % MOD;
domul(x ,s1);
return x;
}
if(x->l == x->r){
LL tmps2 = s2;
s1 = (s1 + y->sum) % MOD;
s2 = (s2 + x->sum) % MOD;
domul(x ,s1);
x->sum = (x->sum + y->sum * tmps2) % MOD;
}
else{
pushdown(x) ,pushdown(y);
x->lch = merge(x->lch ,y->lch ,s1 ,s2);
x->rch = merge(x->rch ,y->rch ,s1 ,s2);
pushup(x);
}
return x;
}
int dep[MX];
void dfs(int x ,int f ,int depth){
dep[x] = depth;
for(int i = head[x] ,d ; i ; i = h[i].next){
if((d = h[i].node) == f) continue;
dfs(d ,x ,depth + 1);
}
root[x] = newnode(0 ,n);
int mx = 0;
for(auto i : limit[x]){
if(dep[i] > dep[x]) continue;
// change(root[x] ,dep[i] ,1);
mx = std::max(mx ,dep[i]);
}
change(root[x] ,mx ,1);
for(int i = head[x] ,d ; i ; i = h[i].next){
if((d = h[i].node) == f) continue;
LL tmp = sum(root[d] ,0 ,dep[x]) ,tmp2 = 0;
merge(root[x] ,root[d] ,tmp ,tmp2);
}
// debug("sum dp[%d] = %lld\n" ,x ,sum(root[x] ,0 ,n));
}
int main(){
n = read();
for(int i = 1 ,u ,v ; i < n ; ++i){
u = read() ,v = read();
addedge(u ,v);
}
int m = read();
for(int i = 1 ,u ,v ; i <= m ; ++i){
u = read() ,v = read();
limit[u].push_back(v);
limit[v].push_back(u);
}
dfs(1 ,0 ,1);
printf("%lld\n" ,sum(root[1] ,0 ,0));
return 0;
}
