P10421 [蓝桥杯 2023 国 A] 树上的路径 题解

· · 题解

P10421 [蓝桥杯 2023 国 A] 树上的路径 题解

题目传送门。

居然没有题解？

本题解为 O(n\log^2n)，拜谢 O(n\log n) 大佬（不知道存不存在单 \log 的做法，存在的话麻烦告知我）。

题意

给一棵树，求 \sum\limits_{i=1}^n{\sum\limits_{j=i+1}^{n}{dis(i,j)\cdot[L \le dis(i,j) \le R]}}。

思路

看到题首先想到了这个题：Tree。

我们在普通点分治的过程中，有一种写法是用一个桶维护深度出现的次数。

在 Tree 这个题中，我们可以把桶改成值域树状数组，修改时 +1，查询前缀和即可。

本题中，我们需要在修改树状数组时加上深度，而不是加上 1，这样就维护出了当前根到儿子的子树中所有节点的深度前缀和。

由于点分治每一层是解决经过根节点的路径，我们在遍历的时候还需要在前面深度前缀和的基础上，加上【当前点到根节点的长度】乘上【前面子树中长度符合长度限制的点的个数】，贡献进答案中。

所以要开两个树状数组，第一个修改时加深度，第二个修改时加 1。

AC 代码

#include<bits/stdc++.h>
#define Code using
#define by namespace
#define wjb std
Code by wjb;
#define int long long
int in()
{
    int k=0,f=1;
    char c=getchar();
    while(c<'0'||c>'9')
    {
        if(c=='-')f=-1;
        c=getchar();
    }
    while(c>='0'&&c<='9')k=k*10+c-'0',c=getchar();
    return k*f;
}
void out(int x)
{
    if(x<0)putchar('-'),x=-x;
    if(x<10)putchar(x+'0');
    else out(x/10),putchar(x%10+'0');
}
const int N=1e6+10;
vector<int>e[N];
int maxp[N],siz[N];
int root;
bool vis[N];
void get_root(int u,int fa,int tot)
{
    siz[u]=1,maxp[u]=0;
    for(int v:e[u])
    {
        if(v==fa||vis[v])continue;
        get_root(v,u,tot);
        siz[u]+=siz[v],maxp[u]=max(maxp[u],siz[v]);
    }
    maxp[u]=max(maxp[u],tot-siz[u]);
    if(!root||maxp[u]<maxp[root])root=u;
}
int stk[N],tot=0;
void dfs(int u,int fa,int d)
{
    stk[++tot]=d;
    for(int v:e[u])
    {
        if(vis[v]||v==fa)continue;
        dfs(v,u,d+1);
    }
}
int n,l,r,ans=0;
struct bit
{
    int c[N];
    void add(int x,int v){for(x++;x<=n+1;x+=x&-x)c[x]+=v;}
    int ask(int x){int ss=0;for(x++;x;x-=x&-x)ss+=c[x];return ss;}
    int ask(int l,int r){return ask(r)-ask(l-1);}
}T,num;
void work(int u)
{
    num.add(0,1),tot=0;int i=1;
    for(int v:e[u])
    {
        if(vis[v])continue;
        dfs(v,u,1);
        for(int j=i;j<=tot;j++)
        {
            int ll=max(0ll,l-stk[j]),rr=r-stk[j];
            if(rr<0)continue;
            ans+=T.ask(ll,rr)+num.ask(ll,rr)*stk[j];
        }
        for(;i<=tot;i++)T.add(stk[i],stk[i]),num.add(stk[i],1);
    }
    for(;tot;tot--)T.add(stk[tot],-stk[tot]),num.add(stk[tot],-1);
    num.add(0,-1);
}
void solve(int u)
{
    vis[u]=true,work(u);
    for(int v:e[u])
    {
        if(vis[v])continue;
        root=0,get_root(v,0,siz[u]);
        solve(root);
    }
}
int divid()
{
    memset(vis,0,sizeof(vis));
    ans=root=0,maxp[0]=n,get_root(1,0,n);
    solve(root);
    return ans;
}
signed main()
{
    n=in(),l=in(),r=in();
    for(int i=2;i<=n;i++)
    {
        int u=in();
        e[u].push_back(i),e[i].push_back(u);
    }
    out(divid());
    return 0;
}

有点卡常，注意一下常数。

注意：十年 oi 一场空，不开 long long 见祖宗！