[hdu5405] sometimes naive——树链剖分+线段树

[hdu5405]sometimes naive

Problem Description

Rhason Cheung had a naive problem, and asked Teacher Mai for help. But Teacher Mai thought this problem was too simple, sometimes naive. So she ask you for help.

She has a tree with n vertices, numbered from 1 to n. The weight of i-th node is wi.

You need to support two kinds of operations: modification and query.

For a modification operation u,w, you need to change the weight of u-th node into w.

For a query operation u,v, you should output ∑Ni=1∑Nj=1 ∑ i = 1 N ∑ j = 1 N . If there is a vertex on the path from u to v and the path from i to j in the tree, f(i,j)=wiwj, otherwise f(i,j)=0. The number can be large, so print the number modulo 109+7

Input

There are multiple test cases.

For each test case, the first line contains two numbers n,m(1≤n,m≤105).

There are n numbers in the next line, the i-th means wi(0≤wi≤109).

Next n−1 lines contain two numbers each, ui and vi, that means that there is an edge between ui and vi.

The following are m lines. Each line indicates an operation, and the format is “1 u w”(modification) or “2 u v”(query)(0≤w≤109)

Output

For each test case, print the answer for each query operation.

Sample Input

6 5
1 2 3 4 5 6
1 2
1 3
2 4
2 5
4 6
2 3 5
1 5 6
2 2 3
1 1 7
2 2 4

Sample Output

341
348
612

题目大意：

给定一颗树，有两种操作，第一种是把点u的权值改成w，第二种操作是求 ∑Ni=1∑Nj=1 ∑ i = 1 N ∑ j = 1 N （满足i到j的路径与u，v有交点）w[i]*w[j]。

大致思路：

对于第二种操作，我们把u到v的路径拉直，展开成为一条链

然后我们可以发现，要求i到j的路径与u到v的路径有交点，唯一不满足的情况就是i和j都在同一颗子树之内，例如都在S2这颗子树之内的时候，i和j的路径必然树不过u到v的路径的。

其他的所有情况都是满足有交点的，所以我们只要把所有的情况算出来减去各个子树内的权值的和的平方就可以了。

所以可以先剖一下，剖成线性结构去维护，然后开两颗线段树，其中一颗sum维护权值和，另外一颗s维护该点的所有轻链子树的平方的和的和。这么一来查询的时候就可以一直沿着重链往上走。
如果u有重儿子的话，显然重儿子这个子树也不在u到v的路径上（因为是从一条轻链上面跳上来的），所以要减去这个重儿子的子树内权值得和的平方。
接下来每次查询一下[in[top[u]],in[u]]区间范围内的s的值的和就可以了，另外由于这一条重链的top必然是它父亲的轻儿子，所以要先加上这条链top的子树的权值和的平方，以抵消top的父亲减去top内的子树的平方和。
其实更新的操作也是类似的，当我们更新了一个点的权值之后，必然有些节点的S值会发生改变，不难发现这些节点全部都是u到根的路径上的每一条的链的最下面那个节点（因为只有这些节点才会有包括了u节点在内的轻儿子），所以这么一来只需要一条一条的往上爬，每爬到一条链更新一下就可以了。
具体的代码如下

/*===============================
 * Auther: ylsoi
 * Problem: hdu5405
 * Time: 2018.1.12
 * ============================*/
#include
#include
#include
#include
#include
#define REP(i,a,b) for(register int i=a;i<=b;++i)
#define MREP(i,x) for(register int i=beg[x];i;i=E[i].last)
#define mid ((l+r)>>1)
#define lson rt<<1,l,mid
#define rson rt<<1|1,mid+1,r
#define ll long long
#define mem(a) memset(a,0,sizeof(a))
using namespace std;
const int maxn=1e5+10;
const ll mod=1e9+7;
int n,m,len,beg[maxn];
struct edge{
    int to;
    int last;
}E[maxn*2];
void add(int u,int v){
    ++len;
    E[len].to=v;
    E[len].last=beg[u];
    beg[u]=len;
}
int son[maxn],size[maxn],fa[maxn],top[maxn],dep[maxn],in[maxn],tot;
ll sum[maxn*4],s[maxn*4],va[maxn],w[maxn];
void dfs1(int u,int f){
    dep[u]=dep[f]+1;
    fa[u]=f;
    size[u]=1;
    son[u]=0;
    MREP(i,u){
        int v=E[i].to;
        if(v==f)continue;
        dfs1(v,u);
        size[u]+=size[v];
        if(size[v]>size[son[u]])
            son[u]=v;
    }
}
void dfs2(int u,int t){
    in[u]=++tot;
    va[tot]=w[u]%mod;
    top[u]=t;
    if(son[u])dfs2(son[u],t);
    MREP(i,u){
        int v=E[i].to;
        if(v==fa[u] || v==son[u])continue;
        dfs2(v,v);
    }
}
void build(int rt,int l,int r){
    if(l==r)sum[rt]=va[l];
    else{
        build(lson);
        build(rson);
        sum[rt]=(sum[rt<<1]+sum[rt<<1|1])%mod;
    }
}
void update(int rt,int l,int r,int pos,ll x,ll *a){
    if(l==r)a[rt]=x;
    else{
        if(pos<=mid)update(lson,pos,x,a);
        else update(rson,pos,x,a);
        a[rt]=(a[rt<<1]+a[rt<<1|1])%mod;
    }
}
ll query(int rt,int l,int r,int L,int R,ll *a){
    if(L>R)return 0;
    ll ret=0;
    if(L<=l && r<=R)return a[rt];
    else{
        if(L<=mid)ret=(ret+query(lson,L,R,a))%mod;
        if(R>=mid+1)ret=(ret+query(rson,L,R,a))%mod;
    }
    return (ret%mod+mod)%mod;
}
void builds(){
    REP(i,1,n){
        int u=i;
        ll SUM=0ll;
        MREP(j,u){
            int v=E[j].to;
            if(v==son[u] || v==fa[u])continue;
            ll qu=query(1,1,n,in[v],in[v]+size[v]-1,sum);
            SUM=(SUM+qu*qu)%mod;
        }
        update(1,1,n,in[u],SUM,s);
    }
}
void update_node(int x,ll value){
    ll add=w[x]-value;
    int v=x;
    v=top[v];
    while(fa[v]){
        int u=fa[v];
        ll SUM=query(1,1,n,in[v],in[v]+size[v]-1,sum);
        ll QU=query(1,1,n,in[u],in[u],s);
        SUM=(QU-(SUM*SUM)+(SUM-add)*(SUM-add))%mod;
        update(1,1,n,in[u],SUM,s);
        v=top[fa[v]];
    }
    w[x]=value;
    update(1,1,n,in[x],value,sum);
}
void cal(int u,int v){
    ll ans=sum[1]*sum[1];
    while(top[u]!=top[v]){
        if(dep[top[u]]if(son[u]){
            int ch=son[u];
            ll qu=query(1,1,n,in[ch],in[ch]+size[ch]-1,sum);
            ans=(ans-qu*qu)%mod;
        }
        ans=(ans-query(1,1,n,in[top[u]],in[u],s))%mod;
        ll qu=query(1,1,n,in[top[u]],in[top[u]]+size[top[u]]-1,sum);
        ans=(ans+qu*qu)%mod;
        u=fa[top[u]];
    }
    if(in[u]>in[v])swap(u,v);
    ans=(ans-query(1,1,n,in[u],in[v],s))%mod;
    ll tot_sum=query(1,1,n,in[u],in[u]+size[u]-1,sum);
    ans=(ans-(sum[1]-tot_sum)*(sum[1]-tot_sum))%mod;
    if(son[v]){
        int ch=son[v];
        ll qu=query(1,1,n,in[ch],in[ch]+size[ch]-1,sum);
        ans=(ans-qu*qu)%mod;
    }
    printf("%lld\n",(ans%mod+mod)%mod);
}
void clear(){
    len=tot=0;
    mem(beg);
    mem(sum);
    mem(s);
}
int main(){
#ifndef ONLINE_JUDGE
    freopen("hdu5405.in","r",stdin);
    freopen("hdu5405.out","w",stdout);
#endif
    while(~scanf("%d%d",&n,&m)){
        REP(i,1,n)scanf("%lld",&w[i]);
        REP(i,1,n-1){
            int u,v;
            scanf("%d%d",&u,&v);
            add(u,v);
            add(v,u);
        }
        dfs1(1,0);
        dfs2(1,1);
        build(1,1,n);
        builds();
        REP(i,1,m){
            int key,u,v;
            scanf("%d%d%d",&key,&u,&v);
            if(key==1)update_node(u,v);
            else cal(u,v);
        }
        clear();
    }
    return 0;
}

总的来说这题还是聚集了几个很不错的算法，同时也很考察代码能力，算是一道十分不错的题目，学习过线段树和树链剖分的同学可以尝试，可能会花上你一点时间，但是做完了之后是一定会有很多的收获的。
这一题的提交网址http://acm.hdu.edu.cn/showproblem.php?pid=5405