一种利用重链剖分优化一类树形动态规划空间复杂度的方法

Origin

在某一场GDOI模拟赛上，一道好好的点分治题目，本蒟蒻强行大力优化暴力碾了过去。这题中我算法的瓶颈不在时间复杂度，而在于空间复杂度。为了解决这个问题，我想到了一种使用重链剖分来优化空间复杂度的方法。

---

Problem

首先我们要明确一下这种方法的适用范围：
一个点对该点父亲的贡献，可以直接利用该点已知的信息以及该点父亲已有的信息计算得，不需要依赖该点的兄弟或者其它点的信息。
比如我们现在要讨论的这一道题目：
给定一棵n个节点的树，每个节点有一个颜色（颜色编号是在[1,k]内的整数，在这题里面k≤10），你需要统计包含所有颜色的树路径数目。
这个显然如果我们处理出每个点为顶点然后颜色集合二进制状态为S的树链条数，就可以直接求出来。现在关键在于，空间限制不能让我对每一个点都开大小为2k的数组，但是这里父亲的信息必须由儿子得到。
怎么压空间呢？考虑这样一种算法：我们建一个内存池动态分配空间（每次会给一个节点分配大小为2k的空间），一开始我们一个点也不分配，然后我们开始深搜这棵树，如果我们遇到了一个叶子节点，就给分配空间记录一下对应的信息。然后在我们退出一个点的时候，先看看父亲是否已经分配了空间。如果它父亲没有被分配空间，我们就先分配一个。然后我们用这个点的信息更新其父亲的信息，并且把分配给这个点的空间释放回内存池。
这样的话，我的空间复杂度是多少呢？我们假定在DFS过程中，一个点第一个访问的儿子叫做偏爱儿子，那么可以发现，这个算法的空间复杂度是这个点到根的路径上，不是父亲的偏爱儿子的点的个数。也就是说，我DFS过程中对儿子的选择，决定了我空间复杂度的多少。
仔细看看这个“偏爱儿子”的定义，就可以发现，这个空间复杂度其实相当于我用某一种方式对这棵树进行树链剖分，然后一个点到根路径上轻边的条数。如果我使用的算法是重链剖分，那么轻边的条数就是logn的，空间便自然可以做到O(2klogn)。

---

Implement

#include <iostream>#include <cstdio>#include <cctype>#include <queue>using namespace std;typedef long long LL;int read(){    int x=0,f=1;    char ch=getchar();    while (!isdigit(ch)) f=ch=='-'?-1:f,ch=getchar();    while (isdigit(ch)) x=x*10+ch-'0',ch=getchar();    return x*f;}const int N=50050;const int V=17;const int K=10;const int S=1<<K;const int E=N<<1;int fid[N],fa[N],size[N],last[N],col[N],lgc[N];int nxt[E],tov[E];queue<int> avail;int f[V][S];int g[S];int n,k,s,tot;LL ans;void insert(int x,int y){tov[++tot]=y,nxt[tot]=last[x],last[x]=tot;}void dfs(int x){    size[x]=1;    for (int i=last[x],y;i;i=nxt[i])        if ((y=tov[i])!=fa[x])        {            fa[y]=x,dfs(y),size[x]+=size[y];            if (!lgc[x]||size[lgc[x]]<size[y]) lgc[x]=y;        }}void pre(){for (int i=0;i<V;++i) avail.push(i);}void makeroom(int x){fid[x]=avail.front(),avail.pop();}void clear(int x){    int id=fid[x];    for (int s_=0;s_<s;++s_) f[id][s_]=0;    fid[x]=-1,avail.push(id);}void process(int id){    for (int s_=0;s_<s;++s_) g[s_]=f[id][s_];    for (int i=0,l,len;i<k;++i)    {        len=1<<i+1,l=len>>1;        for (int j=0;j<s;j+=len)            for (int x=j;x<j+l;++x)                g[x]+=g[x+l];    }}void calc(int x){    int id;    fid[x]=-1;    if (!lgc[x]) makeroom(x);    else    {        calc(lgc[x]);        for (int i=last[x],y;i;i=nxt[i])            if ((y=tov[i])!=fa[x]&&y!=lgc[x]) calc(y);    }    ++f[id=fid[x]][1<<col[x]],process(id);    for (int s_=0;s_<s;++s_) ans+=1ll*f[id][s_]*g[(s-1)^s_];    if (fa[x])    {        int y=fa[x];        if (fid[y]==-1) makeroom(y);        int id_=fid[y],c=1<<col[y];        for (int s_=0;s_<s;++s_) f[id_][s_|c]+=f[id][s_],ans-=1ll*f[id][s_]*g[(s-1)^(s_|c)];    }    clear(x);}int main(){    freopen("colortree.in","r",stdin),freopen("colortree.out","w",stdout);    n=read(),k=read(),s=1<<k;    for (int i=1;i<=n;++i) col[i]=read()-1;    for (int i=1,x,y;i<n;++i) x=read(),y=read(),insert(x,y),insert(y,x);    fa[1]=0,dfs(1),pre(),calc(1);    printf("%lld\n",ans);    fclose(stdin),fclose(stdout);    return 0;}

---

Postscript

这个算法是我在比赛中的脑洞成果，其实我个人认为其应用范围恐怕很狭窄，毕竟对于这一类题目，如果我树形dp的空间都会被卡的话，这个时间复杂度估计也特别大了。一般情况下是不会用到的。不过作为一个有趣的发现，就在这里记录一下吧。
如果各位dalao对这个算法有什么更深入的想法，欢迎留言交流。

---

Update

发现了一道使用这个思路的题目：[HDU5511]Minimum Cut-Cut