2023牛客暑期多校训练营8 D.Distance on Tree

题目大意

给定一颗大小为 $n$ ，根节点为 $1$ 的树，和 $n$ 个点的点权 $a_i$， $n-1$ 条边的边权 $w_i$，和一个值 $m$。有 $q$ 个询问，每次给定两个数 $x_i,k_i$ ，在 $x_i$ 的子树下求三个节点 $u,v,w$ ，满足 $a_u+a_v+a_w\equiv k_i(modm)$，且 $dis(u,v)+dis(u,w)+dis(v,w)$ 的值尽可能大。
定义 $dis(a,b)$ 表示节点 $a$ 和 $b$ 的最短路径边权之和。
$1\le n,m\le 2000,1\le q,w_i\le 2\ast 10^5,0\le k_i,a_i<m$

题解

参考 walk alone 的新手向题解==
首先，它的询问数非常之大，甚至不能 $O(n)$ 遍历完。
所以我们考虑离线的操作，
计算一个 $x_i$ 时往往会对不同的 $k_i$ 做出贡献，所以我们可以将所有的答案全都求出，再进行输出。
显然的，一个节点的状态可以由它的子节点推过来，所以我们选择用树形DP求解。
设 $d{p}_{i,1/2/3,j}$ 表示在节点 $i$ 的子树内，$1/2/3$ 个子节点点权之和模 $m$ 后结果为 $j$ 的到 $i$ 距离的两倍（方便写定义式）。
显然的有
$d{p}_{i,1,j}=max(d{p}_{son,1,j}+2\ast w_{i-son})$
$d{p}_{i,2,j}=max(d{p}_{son,2,j}+2\ast w_{i-son},d{p}_{pre\_son,1,a}+d{p}_{son,1,j-a}+2\ast w_{i-son},d{p}_{i,2,j})$
$d{p}_{i,3,j}=max(d{p}_{son,3,j},d{p}_{i,3,j},d{p}_{pre\_son,1,a}+d{p}_{son,2,r-a}+2\ast w_{i-son},d{p}_{pre\_son,2,a}+d{p}_{son,1,r-a}+2\ast w_{i-son})$
我们发现，如果直接处理子树合并，它的复杂度会到达 $O(n{m}^2)$
讨论：
合并子树时，有许多 $d{p}_{i,1/2/3,j}$ 的 $j$ 范围没有到达 $m$ 在维护这些点时，可以直接枚举子节点中有值的合并转移。
以 $\sqrt{m}$ 为分界线，
定义子树节点 $<\sqrt{m}$ 的为小点，反之为大点，
小点和大点合并，复杂度为 $m$
大点和大点合并，复杂度为 $m^2$，
在一棵树中，大点的个数不大于 $\frac{n}{\sqrt{m}}$，
小点和大点合并，复杂度为 $n\sqrt{m}$
所以总复杂度为 $n{m}^{\frac{3}{2}}$
故复杂度合理。

参考代码

#include<bits/stdc++.h>
using namespace std;
const int N=2e3+5;
struct node{
    int to,w;
};
vector<node> v[N];
int dp[N][4][N];
int l1[N],l2[N];
int f[N],a[N],sum[N];
int n,m,q;
int ans=0,t1,t2;
void merge(int dp[],int a[],int b[],int w) //合并子节点               
{
    t1=0,t2=0;
    for(int i=0;i<m;i++)    //只找有值的维护
    {
        if(a[i]>=0)
            l1[++t1]=i;
        if(b[i]>=0)
            l2[++t2]=i;
    }
    for(int i=1;i<=t1;i++)
        for(int j=1;j<=t2;j++)
            dp[(l1[i]+l2[j])%m]=max(dp[(l1[i]+l2[j])%m],a[l1[i]]+b[l2[j]]+2*w);
}
void dfs(int x,int fa)
{
    dp[x][1][f[x]]=0;
    for(int i=0;i<v[x].size();i++)
    {
        node son=v[x][i];
        if(son.to==fa)
            continue;
        dfs(son.to,x);
        for(int i=3;i>=1;i--)
        {
            for(int j=1;j<i;j++)
                merge(dp[x][i],dp[son.to][j],dp[x][i-j],son.w);
            for(int j=0;j<m;j++)
            {
                if(dp[son.to][i][j]!=-1)
                    dp[x][i][j]=max(dp[x][i][j],dp[son.to][i][j]+(i<=2)*(2*son.w));
//                 cout <<dp[x][i][j]<<" ";
            }
//             puts("");
        }
    }           
}
int main()
{
    scanf("%d%d%d",&n,&m,&q);
    for(int i=1;i<=n;i++)
        scanf("%d",&f[i]);
    for(int i=1;i<n;i++)
    {
        int a,b,c;
        scanf("%d%d%d",&a,&b,&c);
        v[a].push_back({b,c});
        v[b].push_back({a,c});
    }
    for(int i=1;i<=n;i++)
        for(int j=1;j<=3;j++)
            for(int k=0;k<m;k++)
                dp[i][j][k]=-1;
    dfs(1,0);
    while(q--)
    {
        int a,b;
        scanf("%d%d",&a,&b);
        printf("%d\n",max(0,dp[a][3][b]));
    }
}