题解：P17145 [NOI 2026] 彩虹树

Daylighten · 2026-11-11 12:37:47 · 题解

非常好的计数题，给了我大大的震撼。

这道题，我认为最难的地方在于写出一个多项式时间复杂度做法，后面的优化或许没有那么难。

思路

对于这种题，我们可以尝试的方向就是找给定一个缤纷序列 \{c_1,c_2,...,c_n\} 和集合 S，如何判断是否存在一种合法的染色方式。

1. 找判定方法

以下默认 v\in son_u。

首先需要满足的就是 \max c_v\le c_u\le 1+\sum c_v。

然后我们记录一个数组 d_u 表示 u 子树内最少还剩下多少种颜色需要 u 的祖先（不包括 u）来满足同色限制。因为我们想判断是否存在一种合法方案，而 d_u 越小，对祖先的限制就越弱。如果限制最弱的方案都不合法，那么一定不存在合法解了。所以我们需要让 d_u 都尽可能小。最后只用判断 d_1 是否等于 0 即可。

考虑如何对于一个点 u，其所有儿子的 c,d 都知道了，而且 c_u 也是确定的，那么 d_u 最小是多少？

我们先算出 u 没有造成任何影响，只是合并所有儿子的子树后的值 d'。

我们可以把属于不同的子树的两个限制定为同一种颜色，然后将这两个限制合并成一个限制。

当 c_u\le \sum c_v 时，我们最多可以合并 \sum c_v-c_u 种颜色，也就是最多可以合并这么多的限制，剩下 \sum d_v-\sum c_v+c_u 个限制。但是又因为 d'\ge \max d_v，所以 d' 可以写成

d'=\max(\max d_v,\sum d_v-\sum c_v+c_u)

再考虑 u 的对子树内的影响：

首先 u 对子树内可以消除最多一个限制，所以 d' 如果为正就会减一。

当然，特殊的情况是 c_u=\sum c_v+1，此时 u 必须是单独一种颜色。我们可以先按照上面同样的方式算一遍，发现算出来的 d' 恰好是 \sum d_v，是正确的，所以不用特判，把刚刚的 c_u\le \sum c_v 的限制删掉就行了。

接着考虑 u 是否会添加限制：

如果 u\not\in S，此时 u 不会产生新的限制，综上，此时
d_u=\max(0,d'-1)
如果 u\in S，此时 u 会产生新的限制，于是就要从上种情况加一，就是：
d_u=\max(1,d')

综上，我们可以 O(n) 判断一对 (\{c_1,c_2,...,c_n\},S) 是否有贡献。

2. 得出暴力 dp 做法

我们得到了判定条件和流程，接下来就可以考虑 dp 了。

设 dp_{u,c,d} 表示考虑完了 u 及其子树，且 c_u=c,d_u=d 的方案数。

这个过程，就很像 dp 套 dp。

转移的时候，最暴力的做法就是先枚举 c_u，然后算出 d_u 需要 \sum d_v,\sum c_v,\max d_v。同时需要保证 c_v\le c_u。于是直接设 f_{i,j,k} 表示 u 已经考虑过的儿子的 \sum c_v=i,\sum d_v=j,\max d_v=k 的方案数。然后这个就是一个树上背包。

转移就是枚举儿子 v，还有枚举 1\le c_v\le \min(sz_v,c_u),d_v\le c_v，然后更新：

f_{i+c_v,j+d_v,\max(k,d_v)}\gets dp_{v,c_v,d_v}\times f_{i,j,k}

算出 f_{i,j,k} 后，就算出来 dp_u 的值，先算出 d'=\max(k,j-i+c_u)，为了方便理解，这里我们带入 i=\sum c_v,j=\sum d_v,k=\max d_v，那么算出 d'=\max(\max d_v,\sum d_v-\sum c_v+c_u)，然后就可以讨论 u 是否属于 S 了：

\\dp_{u,c_u,\max(1,d')}\gets f_{i,j,k}

这样做时间复杂度是 O(n^6) 的。至少我们得到了一个多项式时间复杂度做法了。

下面放了部分代码，其中变量 sc 和 sd 分别代表 \sum c_v 和 \sum d_v，md 代表 \max d_v，外层枚举的 j 是枚举的 c_u 的值。

:::info[代码实现关键部分]

//add(x,y) 是模意义下 x+=y
for(int j=1;j<=sz[u];j++)
{
    f[0][0][0]=1;
    int now=0,mxd=0;
    for(auto v:e[u])
    {
        for(int sc=0;sc<=now;sc++)
        {
            for(int sd=0;sd<=sc;sd++)
            {
                for(int md=0;md<=mxd;md++)
                {
                    if(!f[sc][sd][md]) continue;
                    for(int c=1;c<=min(j,sz[v]);c++)
                    {
                        for(int d=0;d<=c;d++)
                        {
                            if(!dp[v][c][d]) continue;
                            add(g[sc+c][sd+d][max(md,d)],1ll*f[sc][sd][md]*dp[v][c][d]%mod);
                        }
                    }
                }
            }
        }
        mxd=max(mxd,sz[v]);
        now+=min(j,sz[v]);
        for(int sc=0;sc<=now;sc++)
        {
            for(int sd=0;sd<=sc;sd++)
            {
                for(int md=0;md<=mxd;md++) f[sc][sd][md]=g[sc][sd][md],g[sc][sd][md]=0;
            }
        }
    }
    for(int sc=j-1;sc<=now;sc++)
    {
        for(int sd=0;sd<=sc;sd++)
        {
            for(int md=0;md<=mxd;md++)
            {
                int nk=max(md,sd-(sc-j));
                add(dp[u][j][max(0,nk-1)],f[sc][sd][md]);
                add(dp[u][j][max(1,nk)],f[sc][sd][md]);
            }
        }
    }
    for(int sc=0;sc<=now;sc++) for(int sd=0;sd<=sc;sd++) for(int md=0;md<=mxd;md++) f[sc][sd][md]=0;
}

:::

3. 优化

接下来就是优化了。

发现外层枚举 c_u 很慢，因为每个不同的 c_u 我们都需要跑一遍树上背包。

注意一下 d'=\max(\max d_v,\sum d_v-\sum c_v+c_u) 和 c_u 的关系。可以发现 d' 关于 c_u 的函数是一个水平线加上一条斜率为一的斜线构成的分段函数，分割点是 T=\max d_v+\sum c_v-\sum d_v。

我们分别考虑这两段：

对于左边那个水平段，其 c_u 的定义域是 [\max c_v,T]。值都是 \max d_v。

我们还是需要维护至少 \max d_v,\max c_v,\sum {c_v}-\sum d_v 三个信息才能算，但是因为我们更新的是一段连续的 c_u，所以考虑差分。

此时发现差分的左右两个端点都只用维护两个信息就可以算出来：
- 对于左边那个端点，c_u=\max c_v,d'=\max d_v。只用知道 \max c_v 和 \max d_v 两个信息就可以表示出来。
- 对于右边那个端点，c_u=T+1,d'=\max d_v。只用知道 \max d_v 和 \sum c_v-\sum d_v就可以表示出来。
于是可以用上面算 f 的方式，我们可以设一个 g_{i,j} 为 \max c_v=i,\max d_v=j 的方案数，和 h_{i,j} 为 \max d_v=i,\sum c_v-\sum d_v=j 的方案数。这样状态数就可以少一维算出 g,h，然后开一个差分数组，存下所有更新。

这里给出 g 的计算代码，下面用的变量名是 f：

:::info[部分代码]
```
f[0][0]=1;
int mx=0,now=0;// max(c), max(d)
for(auto v:e[u])
{
    for(int mc=0;mc<=mx;mc++)
    {
        for(int md=0;md<=mx;md++)
        {
            if(!f[mc][md]) continue;
            for(int c=1;c<=sz[v];c++)
            {
                for(int d=0;d<=c;d++)
                {
                    if(!dp[v][c][d]) continue;
                    add(tmp[max(mc,c)][max(md,d)],1ll*f[mc][md]*dp[v][c][d]%mod);
                }
            }
        }
    }
    mx=max(mx,sz[v]);
    now+=sz[v];
    for(int i=0;i<=mx;i++)
    {
        for(int j=0;j<=mx;j++) f[i][j]=tmp[i][j],tmp[i][j]=0;
    }
}
for(int mc=0;mc<=mx;mc++)
{
    for(int md=0;md<=mx;md++)
    {
        if(!f[mc][md]) continue;
        add(cha1[mc][max(0,md-1)],f[mc][md]);
        add(cha1[mc][max(1,md)],f[mc][md]);
        f[mc][md]=0;
    }
}
```
:::
对于斜率为一的右边那一段，定义域为 [T+1,\sum c_v+1]，值域 [\max d_v+1,\sum d_v+1]

同样地，可以注意到我们同样可以使用差分，而且两个端点都只用两个信息就可以知道了：
- 左端点 c_u=T+1,d'=\max d_v+1，只用知道 \max d_v 和 \sum c_v-\sum d_v 就可以了。
- 右端点 c_u=\sum c_v+2,d'=\sum d_v+2，只用知道 \sum c_v 和 \sum d_v 就可以了。
而且注意到左端点需要的信息在上面一种情况已经算了一遍了，因为本质上是一个点。所以只用再跑一遍记录 \sum c_v,\sum d_v 的树上背包就可以了。然后差分是斜着的差分。

综上，我们只用跑三遍记录两个信息的树上背包，再加上两次差分，就可以算出 dp_u。

于是时间复杂度优化到了 O(n^4)，可以通过。

代码比较丑，因为我把三次树上背包复制了三遍分开写的。

:::info[代码]

#include<bits/stdc++.h>
using namespace std;
const int N=205,mod=998244353;
int n,dp[N][N][N],sz[N];
int f[N][N],tmp[N][N];
int cha1[N][N],cha2[N][N];
vector<int> e[N];
int read()
{
    int x=0,f=1;char ch=getchar();
    while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
    while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
    return x*f;
}
void add(int &x,int y)
{
    x+=y;
    if(x>=mod) x-=mod;
}
void dec(int &x,int y)
{
    x-=y;
    if(x<0) x+=mod;
}
void update(int x,int y)
{
    for(int i=0;i<=x;i++)
    {
        for(int j=0;j<=y;j++) f[i][j]=tmp[i][j],tmp[i][j]=0;
    }
}
void dfs(int u,int fa)
{
    sz[u]=1;
    if(!e[u].size())
    {
        dp[u][1][1]=1;
        dp[u][1][0]=1;
        return;
    }
    for(auto v:e[u]) dfs(v,u),sz[u]+=sz[v];
    f[0][0]=1;
    int mx=0,now=0;// max(c), max(d)
    for(auto v:e[u])
    {
        for(int mc=0;mc<=mx;mc++)
        {
            for(int md=0;md<=mx;md++)
            {
                if(!f[mc][md]) continue;
                for(int c=1;c<=sz[v];c++)
                {
                    for(int d=0;d<=c;d++)
                    {
                        if(!dp[v][c][d]) continue;
                        add(tmp[max(mc,c)][max(md,d)],1ll*f[mc][md]*dp[v][c][d]%mod);
                    }
                }
            }
        }
        mx=max(mx,sz[v]);
        now+=sz[v];
        update(mx,mx);
    }
    for(int mc=0;mc<=mx;mc++)
    {
        for(int md=0;md<=mx;md++)
        {
            if(!f[mc][md]) continue;
            add(cha1[mc][max(0,md-1)],f[mc][md]);
            add(cha1[mc][max(1,md)],f[mc][md]);
            f[mc][md]=0;
        }
    }

    f[0][0]=1;
    mx=0,now=0;// sum(c)-sum(d), max(d)
    for(auto v:e[u])
    {
        for(int sum=0;sum<=now;sum++)
        {
            for(int md=0;md<=mx;md++)
            {
                if(!f[sum][md]) continue;
                for(int c=1;c<=sz[v];c++)
                {
                    for(int d=0;d<=c;d++)
                    {
                        if(!dp[v][c][d]) continue;
                        add(tmp[sum+c-d][max(md,d)],1ll*f[sum][md]*dp[v][c][d]%mod);
                    }
                }
            }
        }
        mx=max(mx,sz[v]);
        now+=sz[v];
        update(now,mx);
    }
    for(int sum=0;sum<=now;sum++)
    {
        for(int md=0;md<=mx;md++)
        {
            if(!f[sum][md]) continue;
            int mid=md+sum;
            dec(cha1[mid+1][max(0,md-1)],f[sum][md]);
            dec(cha1[mid+1][max(1,md)],f[sum][md]);
            add(cha2[mid+1][max(0,md)],f[sum][md]);
            add(cha2[mid+1][max(1,md+1)],f[sum][md]);
            f[sum][md]=0;
        }
    }

    f[0][0]=1;
    mx=0,now=0;// sum(c), sum(d)
    for(auto v:e[u])
    {
        for(int sc=0;sc<=now;sc++)
        {
            for(int sd=0;sd<=now;sd++)
            {
                if(!f[sc][sd]) continue;
                for(int c=1;c<=sz[v];c++)
                {
                    for(int d=0;d<=c;d++)
                    {
                        if(!dp[v][c][d]) continue;
                        add(tmp[sc+c][sd+d],1ll*f[sc][sd]*dp[v][c][d]%mod);
                    }
                }
            }
        }
        mx=max(mx,sz[v]);
        now+=sz[v];
        update(now,now);
    }
    for(int sc=0;sc<=now;sc++)
    {
        for(int sd=0;sd<=now;sd++)
        {
            if(!f[sc][sd]) continue;
            dec(cha2[sc+2][max(0,sd+1)],f[sc][sd]);
            dec(cha2[sc+2][max(1,sd+2)],f[sc][sd]);
            f[sc][sd]=0;
        }
    }

    for(int a=0;a<=sz[u];a++)
    {
        for(int b=0;b<=a;b++)
        {
            if(a) add(cha1[a][b],cha1[a-1][b]);
            if(a&&b) add(cha2[a][b],cha2[a-1][b-1]);
            add(dp[u][a][b],cha1[a][b]);
            add(dp[u][a][b],cha2[a][b]);
        }
    }
    for(int a=0;a<=sz[u]+1;a++)
    {
        for(int b=0;b<=sz[u]+1;b++) cha1[a][b]=cha2[a][b]=0;
    }
}
int rainbow(int ID,int N,vector<int> f)
{
    n=N;
    for(int i=2;i<=n;i++) e[f[i-1]+1].push_back(i);
    dfs(1,0);
    int ans=0;
    for(int i=1;i<=n;i++) add(ans,dp[1][i][0]);
    return ans;
}
// int main()
// {
//     int C=read(),n=read();
//     vector<int> f;
//     f.push_back(0);
//     for(int i=2;i<=n;i++)
//     {
//         int x=read();
//         f.push_back(x);
//     }
//     cout<<rainbow(C,n,f);
//     return 0;
// }

:::