Codeforces 1606F - Tree Queries(虚树+树形 dp)

时间:2022-11-11 11:38:16

Codeforces 题面传送门 & 洛谷题面传送门

显然我们选择删除的点连同 \(u\) 会形成一个连通块,否则我们如果选择不删除不与 \(u\) 在同一连通块中的点,答案一定更优。

注意到如果我们选择删除 \(u\) 的某个儿子 \(v\),那么答案的增量为 \(chd_v-1-k\),其中 \(chd_v\) 为节点 \(v\) 儿子的个数。而初始时刻答案为 \(chd_u\) 是个定值,因此我们的任务可以等效于,给每个点赋上一个点权 \(chd_u-1-k\),然后要找到一个以 \(u\) 为根的树上连通块,使得这个连通块中除去点 \(u\) 之外其他点点权之和尽可能大。

因此我们有了一个复杂度 \(\Theta(n^2)\) 的做法:对于每一个 \(k\) 都跑一遍树形 \(dp\),即设 \(dp_u\) 表示以 \(u\) 为根的连通块点权之和的最大值,那么显然有转移 \(dp_u=chd_u-1-k+\sum\limits_{v\in son_u}\max(dp_v,0)\)。注意到我们选出的树上连通块边界上(也就是所有满足 \(x\) 在树上连通块中,但 \(x\) 的所有儿子都不在树上连通块)中的点的点权必定都是正的,否则我们扣掉那些点权非正且在连通块边界上的点,答案肯定变得更优。因此考虑对于每个 \(k\) 将所有点权为正的点拎出来建一棵虚树——由于 \(\sum\limits_{i=1}^nchd_i=n-1\),所以对于所有 \(k\),点权为正的点的总个数是 \(\Theta(n)\) 的,然后对它们跑树形 \(dp\),然后每次查询一个点 \(u\) 为根的连通块点权之和的最大值时,如果点 \(u\) 本身就在虚树上,直接返回 \(u\) 的 DP 值即可,否则我们找到 \(u\) 在虚树上所在的链下方的节点 \(v\)——这个可以拿个 set 存储所有在虚树上的节点的 DFS 序,然后直接在 setlower_bound。那么以 \(u\) 为根的连通块点权和的最大值就是 \(dp_v\) 加上 \(fa_v\to u\) 这段路径和的点权和。

时间复杂度 \(\Theta(n\log n)\)​​。

const int MAXN=2e5;
const int LOG_N=18;
int n,qu,hd[MAXN+5],to[MAXN*2+5],nxt[MAXN*2+5],ec=0;
void adde(int u,int v){to[++ec]=v;nxt[ec]=hd[u];hd[u]=ec;}
int chd[MAXN+5],dep[MAXN+5],fa[MAXN+5][LOG_N+2];
int dfn[MAXN+5],tim=0,rid[MAXN+5],edt[MAXN+5];
void dfs0(int x,int f){
fa[x][0]=f;rid[dfn[x]=++tim]=x;
for(int e=hd[x];e;e=nxt[e]){
int y=to[e];if(y==f) continue;
chd[x]++;dep[y]=dep[x]+1;dfs0(y,x);
} edt[x]=tim;
}
int sum_chd[MAXN+5];
void dfs_chd(int x,int f){
for(int e=hd[x];e;e=nxt[e]){
int y=to[e];if(y==f) continue;
sum_chd[y]=sum_chd[x]+chd[y];
dfs_chd(y,x);
}
}
void lca_init(){
dep[1]=1;dfs0(1,0);
for(int i=1;i<=LOG_N;i++) for(int j=1;j<=n;j++)
fa[j][i]=fa[fa[j][i-1]][i-1];
}
int getlca(int x,int y){
if(dep[x]<dep[y]) swap(x,y);
for(int i=LOG_N;~i;i--) if(dep[x]-(1<<i)>=dep[y]) x=fa[x][i];
if(x==y) return x;
for(int i=LOG_N;~i;i--) if(fa[x][i]!=fa[y][i]) x=fa[x][i],y=fa[y][i];
return fa[x][0];
}
vector<int> pos[MAXN+5];
int stk[MAXN+5],tp=0;set<int> st;
vector<pii> qv[MAXN+5];
ll calc_sum(int u,int v,int k){//u is ancestor of v
return (sum_chd[v]-sum_chd[u])-1ll*(k+1)*(dep[v]-dep[u]);
}
int calc_val(int u,int k){return chd[u]-1-k;}
vector<int> g[MAXN+5];
ll dp[MAXN+5],res[MAXN+5];
void insert(int x){
if(!tp) return stk[++tp]=x,void();
int lc=getlca(x,stk[tp]);
// printf("LCA of %d %d is %d\n",x,stk[tp],lc);
// printf("stack: ");
// for(int i=1;i<=tp;i++) printf("%d%c",stk[i]," \n"[i==tp]);
while(tp>1&&dep[stk[tp-1]]>dep[lc]){
// printf("adde %d %d\n",stk[tp-1],lc);
g[stk[tp-1]].pb(stk[tp]);--tp;
}
if(tp&&dep[stk[tp]]>dep[lc]){
// printf("adde %d %d\n",lc,stk[tp]);
g[lc].pb(stk[tp--]);
}
if(!tp||stk[tp]!=lc) stk[++tp]=lc;
stk[++tp]=x;
}
void fin(){
while(tp>=2){
// printf("adde %d %d\n",stk[tp-1],stk[tp]);
g[stk[tp-1]].pb(stk[tp]);--tp;
} tp=0;
}
void dfs_dp(int x,int k){
dp[x]=calc_val(x,k);st.insert(dfn[x]);
for(int y:g[x]){
dfs_dp(y,k);
dp[x]+=max(dp[y]+calc_sum(x,y,k)-calc_val(y,k),0ll);
// printf("calc_sum %d %d %d = %lld\n",x,y,k,calc_sum(x,y,k));
} //printf("DP %d %lld\n",x,dp[x]);
}
void clr(int x){for(int y:g[x]) clr(y);g[x].clear();}
int main(){
scanf("%d",&n);
for(int i=1,u,v;i<n;i++) scanf("%d%d",&u,&v),adde(u,v),adde(v,u);
lca_init();sum_chd[1]=chd[1];dfs_chd(1,0);scanf("%d",&qu);
// for(int i=1;i<=n;i++) printf("chd[%d]=%d\n",i,chd[i]);
for(int i=1;i<=n;i++) for(int j=0;j<chd[i];j++) pos[j].pb(i);
for(int i=1,u,k;i<=qu;i++) scanf("%d%d",&u,&k),qv[k].pb(mp(u,i));
for(int i=0;i<=MAXN;i++){
st.clear();
if(!pos[i].empty()){
// printf("solving %d\n",i);
sort(pos[i].begin(),pos[i].end(),[&](int x,int y){return dfn[x]<dfn[y];});
if(pos[i][0]!=1) insert(1);
for(int x:pos[i]) insert(x);
fin();dfs_dp(1,i);
}
for(pii p:qv[i]){
int u=p.fi,id=p.se;
set<int>::iterator it=st.lower_bound(dfn[u]);
if(it==st.end()||(*it)>edt[u]) res[id]=chd[u];
else{
int pt=rid[*it];
// printf("%d %d\n",id,pt);
// printf("%lld\n",calc_sum(fa[u][0],fa[pt][0],i));
res[id]=chd[u]+max(0ll,dp[pt]+calc_sum(fa[u][0],fa[pt][0],i)-calc_val(u,i));
}
}
clr(1);
}
for(int i=1;i<=qu;i++) printf("%lld\n",res[i]);
return 0;
}
/*
11
1 2
2 3
2 4
1 5
5 6
5 7
7 8
7 9
7 10
1 11
4
1 0
1 1
5 2
11 0
*/