题解 P7581 「RdOI R2」路径权值(distance)

· · 题解

出题人想让我们写长链剖分，所以我们就不写长链剖分。我们写树上启发式合并。

考虑每个点开一个类似数组的东西，记该点的子树对于每个 dep 的 (sum=\sum dis),tot,ans。

那么合并两个子树时,对于任意两个点 x,y，dis(x,y)=dis_x+dis_y-2dis_{lca}。

设我们合并 u 和其儿子 v 所在子树，则对于每个 dep，

ans_u&=ans_u+ans_v+\sum_{x\in u}\sum_{y\in v}dis(x,y)\\ &=ans_u+ans_v+\sum_{x\in u}\sum_{y\in v}dis_x+dis_y-2dis_u\\ &=ans_u+ans_v+\sum_{x\in u}dis_x\sum_{y\in v}1+\sum_{y\in v}dis_y\sum_{x\in u}1-2dis_u\sum_{x\in u}\sum_{y\in v}1\\ &=ans_u+ans_v+sum_utot_v+sum_vtot_u-2tot_utot_vdis_u \end{aligned}

现在的问题是对于每一个 dep 都进行一次合并。我们选择使用 set，然后启发式合并，时间复杂度 \mathbf O(n\log n)，空间复杂度 \mathbf O(n)。

至于为什么启发式合并 set 是 \mathbf O(n\log n) 的，首先遍历 set 是 \mathbf O(n) 的，然后按顺序插入也是 \mathbf O(n) 的，所以摊下来复杂度就是对的。

而空间复杂度是 \mathbf O(n) 的原因是我们可以每次合并后把那个没用的 set 清空，这样可以保证同一时刻所有 set 中仅有 \mathbf O(n) 个节点。

对于每个询问，把它挂到对应节点上，查这个点 dep+k 的答案就好。

代码

#include<cstdio>
#include<set>
#include<vector>
#include<algorithm>
const int N=1e6+10,logN=30,p=1e9+7;
int n,m,rt[N],ans[N],dep[N],dis[N];
struct edge
{
    int v,w,la;
    inline void set(int _v,int _w,int _la){v=_v;w=_w;la=_la;}
}e[N<<1];
struct node
{
    int val,tot,sum,ans;
    bool operator< (const node &x) const {return val<x.val;}
};
int la[N],cnt;
inline void add(int u,int v,int w){e[++cnt].set(v,w,la[u]);la[u]=cnt;}
inline void add(){int a,b,c;scanf("%d%d%d",&a,&b,&c);add(a,b,c);add(b,a,c);}
std::vector<std::pair<int,int> >qu[N];
std::set<node>c[N];
int merge(int x,int y,int z)
{
    if(c[x].size()<c[y].size())x^=y^=x^=y;
    std::set<node>::iterator it,i;node tmp;
    for(i=c[y].begin();i!=c[y].end();++i)
    {
        it=c[x].find(*i);
        if(it==c[x].end()){c[x].insert(*i);continue;}
        tmp=*it;c[x].erase(it);
        tmp.ans=(tmp.ans+i->ans+1ll*tmp.sum*i->tot%p+1ll*tmp.tot*i->sum%p-2ll*tmp.tot*i->tot%p*z%p+p)%p;
        tmp.sum=(tmp.sum+i->sum)%p;
        tmp.tot=(tmp.tot+i->tot)%p;
        c[x].insert(tmp);
    }
    c[y].clear();return x;
}
void dfs(int u,int f)
{
    dep[u]=dep[f]+1;c[rt[u]=u].insert((node){dep[u],1,dis[u],0});
    for(int i=la[u];i;i=e[i].la)
    if(e[i].v!=f)dis[e[i].v]=(dis[u]+e[i].w)%p,
    dfs(e[i].v,u),rt[u]=merge(rt[u],rt[e[i].v],dis[u]);
    int i,n=qu[u].size();
    std::set<node>::iterator it;
    for(i=0;i<n;i++)
    {
        it=c[rt[u]].find((node){dep[u]+qu[u][i].first,0,0,0});
        if(it!=c[rt[u]].end())ans[qu[u][i].second]=it->ans;
    }
}
int main()
{
    int i,x,y;scanf("%d%d",&n,&m);
    for(i=1;i<n;i++)add();
    for(i=1;i<=m;i++)scanf("%d%d",&x,&y),qu[x].push_back(std::make_pair(y,i));
    dfs(1,0);
    for(i=1;i<=m;i++)printf("%d\n",ans[i]);
    return 0;
}