2018.9 ECNU ICPC/CCPC Trial Round #2 Query On Tree (树链剖分+线段树维护)

传送门:https://acm.ecnu.edu.cn/contest/105/problem/Q/

一棵树,支持两种操作:给一条路径上的节点加上一个等差数列;求两点路径上节点和.

很明显,熟练剖分.用线段树维护链上的区间和,每个节点中记录等差数列的首项,公差和区间和.因为两个等差数列叠加之后还是等差数列,所以将首项与公差视作懒惰标记.

因为在寻找LCA的过程中,u往上跳的时候,其实是要维护递减的等差数列(dfs序是u->topu递减,而数列是u->topu递增);v往上跳的时候,是递增的.计算出u->v上路径的长度,就可以根据等差数列通项公式求出an.

#include<cstdio>
#include<iostream>
#include<algorithm>
#include<cstring>
#define lson rt<<1
#define rson rt<<1|1
#define Lson l,m,lson
#define Rson m+1,r,rson
using namespace std;
typedef long long LL;
const int mod = 1e9+7;
const int maxn =5e4+5;
struct Edge{
    int to,next;
}E[maxn<<1];
int n,head[maxn],tot;
int cnt,idx,siz[maxn],fa[maxn],son[maxn],dep[maxn],top[maxn],id[maxn],rnk[maxn];
int a[maxn];
struct Node{
    LL a1,d,sum;                //区间首项,公差,区间和
}tree[maxn<<2];

LL qpow(LL a,LL N)
{
    LL res=1;
    while(N){
        if(N&1) res = res*a %mod;
        a = a*a%mod;
        N>>=1;
    }
    return res;
}

const LL rev2 = qpow(2,mod-2);              //2逆元

void init()
{
    cnt=idx=tot=0;
    memset(head,-1,sizeof(head));
    dep[1]=0,fa[1]=1,siz[0]=0;
    memset(son,0,sizeof(son));
}

void AddEdge(int u,int v)
{
    E[tot] = (Edge){v,head[u]};
    head[u]=tot++;
}

void dfs1(int u)
{
    siz[u]=1;
    for(int i=head[u];~i;i=E[i].next){
        int v=E[i].to;
        if(v!=fa[u]){
            fa[v]=u;
            dep[v]=dep[u]+1;
            dfs1(v);
            siz[u]+=siz[v];
            if(siz[son[u]]<siz[v]) son[u]=v;
        }
    }
}

void dfs2(int u,int topu)
{
    top[u]= topu;
    id[u] = ++idx;
    rnk[idx] = u;
    if(!son[u]) return;
    dfs2(son[u],top[u]);
    for(int i=head[u];~i;i=E[i].next){
        int v=E[i].to;
        if(v!=fa[u]&&v!=son[u]) dfs2(v,v);
    }
}

//------------------------------------线段树

void pushup(int rt){
    tree[rt].sum = (tree[lson].sum + tree[rson].sum)%mod;
}

void pushdown(int l,int r,int rt){
    if(tree[rt].a1 !=0 || tree[rt].d !=0){
        LL a1 = tree[rt].a1, d = (tree[rt].d+mod)%mod;
        int m = (l+r)>>1;
        int n1 = m-l+1, n2 = r-m;
        LL f1 = a1%mod, f2 = (a1 + n1*d%mod +mod)%mod;
        tree[lson].a1 = (tree[lson].a1+f1)%mod;
        tree[rson].a1 = (tree[rson].a1+f2)%mod;
        tree[lson].d = (tree[lson].d+d+mod)%mod;
        tree[rson].d = (tree[rson].d+d+mod)%mod;
        tree[lson].sum = (tree[lson].sum+ f1*n1 %mod+
                        n1*(n1-1)%mod *d %mod*rev2 %mod + mod)%mod;
        tree[rson].sum = (tree[rson].sum+ f2*n2 %mod+
                        n2*(n2-1)%mod *d %mod *rev2 %mod +mod)%mod;
        //cout<<(lson)<<":"<<tree[lson].sum<<" "<<(rson)<<":"<<tree[rson].sum<<endl;
        tree[rt].a1 = tree[rt].d = 0;
    }
}

void build(int l,int r,int rt)
{
    tree[rt].a1 = tree[rt].d  =0;
    if(l==r){
        tree[rt].sum = 0;
        return;
    }
    int m = (l+r)>>1;
    build(Lson);
    build(Rson);
    pushup(rt);
}

void update(int L,int R,LL a1,LL v,int l=1,int r=n,int rt=1){
    if(L<=l && R>=r){
        int nn = r-l+1;
        v = (v+mod)%mod;
        LL f1 = (a1 + (l-L)*v %mod +mod)%mod;                   //叠加的等差数列首项
        tree[rt].a1 = (tree[rt].a1+f1)%mod;
        tree[rt].d = (tree[rt].d+v+mod)%mod;                //公差
        tree[rt].sum = (tree[rt].sum+f1*nn %mod +nn*(nn-1)%mod
                        *v %mod*rev2%mod + mod)%mod;
        return;
    }
    pushdown(l,r,rt);
    int m =(l+r)>>1;
    if(L<=m) update(L,R,a1,v,Lson);                 //首项对齐
    if(R>m) update(L,R,a1,v,Rson);                  //首项对齐
    pushup(rt);
}

LL query(int L,int R,int l=1,int r=n,int rt=1){        //区间查询
    if(L<=l && R>=r) return tree[rt].sum;
    pushdown(l,r,rt);
    int m = (l+r)>>1,ans=0;
    LL res=0;
    if(L<=m) res = (res+query(L,R,Lson)+mod)%mod;
    if(R>m)  res= (res+query(L,R,Rson)+mod)%mod;
    pushup(rt);
    return res;
}

//----------------------------------树剖
void debug()
{
    for(int i=1;i<=n;++i) printf("%lld ",query(id[i],id[i]));
    cout<<endl;
}

int getdist(int u,int v)
{       //计算两点之间的路径长度
    int ans=0;
    while(top[u]!=top[v]){
        if(dep[top[u]]<dep[top[v]]) swap(u,v);
        ans += id[u] - id[top[u]]+1;
        u = fa[top[u]];
    }
    if(dep[u]>dep[v]) swap(u,v);
    ans += id[v]-id[u];
    return ans;
}

void UPDATE(int u,int v,LL a1, LL w)
{
    int nn = getdist(u,v);
    LL an = (a1 + nn*w)%mod;
    //cout<<nn<<" "<<a1<<" "<<an<<endl;
    while(top[u]!=top[v]){
        if(dep[top[u]]<dep[top[v]]){
            int cnt = id[v] - id[top[v]]+1;
            update(id[top[v]],id[v], (an-(cnt-1)*w %mod+mod)%mod,w);            //递增
            an = (an - (cnt*w%mod) +mod)%mod;
            v = fa[top[v]];
        }
        else{
            int cnt = id[u] - id[top[u]] + 1;
            update(id[top[u]],id[u],(a1+(cnt-1)*w)%mod,(mod-w)%mod);            //反向
            a1 = (a1+ cnt*w %mod)%mod;
            u = fa[top[u]];
        }
        //cout<<a1<<" "<<an<<endl;
    }
    if(dep[u]<dep[v]){
        int cnt = id[v] - id[u]+1;
        update(id[u],id[v],an-(cnt-1)*w,w);            //递增
    }
    else{
        int cnt = id[u] - id[v] +1;
        update(id[v],id[u],a1+(cnt-1)*w,-w);            //反向
    }
}

LL Qsum(int u,int v)
{
    int ans=0;
    while(top[u]!=top[v]){
        if(dep[top[u]]<dep[top[v]]) swap(u,v);
        ans = (ans+query(id[top[u]],id[u])) %mod;
        u = fa[top[u]];
    }
    if(dep[u]>dep[v]) swap(u,v);
    ans = (ans + query(id[u],id[v]))%mod;
    return ans;
}

int main()
{
    #ifndef ONLINE_JUDGE
        freopen("in.txt","r",stdin);
        freopen("out.txt","w",stdout);
    #endif
    int m,q,u,v,op;
    while(scanf("%d%d",&n,&q)==2){
        init();
        for(int i=1;i<n;++i){
            scanf("%d%d",&u,&v); u++,v++;
            AddEdge(u,v);
            AddEdge(v,u);
        }
        dfs1(1);
        dfs2(1,1);
        build(1,n,1);
        LL w;
        while(q--){
            scanf("%d",&op);
            if(op==1){
                scanf("%d %d %lld",&u,&v,&w); u++, v++;
                if(w==0) continue;
                UPDATE(u,v,w,w);
                //debug();
            }
            else{
                scanf("%d %d",&u,&v); u++, v++;
                printf("%lld\n",Qsum(u,v));
            }
        }
    }
    return 0;
}