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

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对这个算法有什么更深入的想法，欢迎留言交流。