HDU 5405 (树链剖分+线段树）

Problem Sometimes Naive

题目大意

　　给你一棵n个节点的树，有点权。

　　要求支持两种操作：

　　　　操作1：更改某个节点的权值。

　　　　操作2：给定u，v， 求 Σw[i][j]   i , j 为任意两点且i到j的路径与u到v的路径相交。

解题分析

　　容易发现对于一个询问，答案为总点权和的平方 减去 去掉u--v这条链后各个子树的点权和的平方的和。 

　　开两棵线段树，tag1记录点权和，tag2记录某点的所有轻链子树的点权和的平方的和。

　　每次沿着重链往上走时，直接加上这条重链的所有点的tag2和，若有重儿子则直接用tag1计算。由于该条重链必定为其父亲的轻链，故为防止计算重复，还需减去该重链所有点的tag1平方和。

　　最后爬到同一颗重链后，还需计算重链上方所有点的贡献。

参考程序

#include <map>
#include <set>
#include <stack>
#include <queue>
#include <cmath>
#include <ctime>
#include <string>
#include <vector>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cassert>
#include <iostream>
#include <algorithm>
#pragma comment(linker,"/STACK:102400000,102400000")
using namespace std;

#define V 100008             
#define E 200008   
#define LL long long
#define lson l,m,rt<<1
#define rson m+1,r,rt<<1|1 
#define clr(x,v) memset(x,v,sizeof(x));
#define rep(x,y,z) for (int x=y;x<=z;x++)
#define repd(x,y,z) for (int x=y;x>=z;x--)
const int mo  = 1000000007;
const int inf = 0x3f3f3f3f;
const int INF = 2000000000;
/**************************************************************************/ 
int n,m,tot;
int val[V];
int size[V],fa[V],w[V],top[V],rk[V],dep[V],son[V];

struct line{
    int u,v,nt;
    line(int u=0,int v=0,int nt=0):u(u),v(v),nt(nt){}
}eg[E];
int lt[V],sum;

void add(int u,int v){
    eg[++sum]=line(u,v,lt[u]); lt[u]=sum;    
}

struct Segment_Tree{
    LL sum[V<<2];
    void clear(){
        clr(sum,0);
    }
    void pushup(int rt){
        sum[rt] = (sum[rt<<1] + sum[rt<<1|1]) % mo;
    }
    void update(int x,int val,int l,int r,int rt){
        if (l==r){
            sum[rt] += val;
            sum[rt] = sum[rt] % mo;
            return;
        }
        int m=(l+r)>>1;
        if (x <= m) update(x,val,lson);
        if (m <  x) update(x,val,rson);
        pushup(rt);
    }
    LL query(int L,int R,int l,int r,int rt){
        if (L<=l && r<=R){
            return sum[rt];
        }
        int m=(l+r)>>1;
        LL res=0;
        if (L <= m) res += query(L,R,lson);
        if (m <  R) res += query(L,R,rson);
        res = res % mo;
        return res;
    }

}Ts,Td;
void init(){
    clr(lt,0); sum=1; tot=0;
    Ts.clear();
    Td.clear();
}
void dfs_1(int u){
    dep[u]=dep[fa[u]]+1; size[u]=1; son[u]=0; 
    for (int i=lt[u];i;i=eg[i].nt){
        int v=eg[i].v;
        if (v==fa[u]) continue;
        fa[v]=u;
        dfs_1(v);
        if (size[v]>size[son[u]]) son[u]=v;
        size[u]+=size[v];
    }
}

void dfs_2(int u,int tp){
    top[u]=tp; w[u]=++tot; rk[tot]=u;;
    if (son[u]) dfs_2(son[u],tp);
    for (int i=lt[u];i;i=eg[i].nt){
        int v=eg[i].v;
        if (v==fa[u]||v==son[u]) continue;
        dfs_2(v,v);
    }
}

int sqr(int x){return 1ll*x*x %mo;}
void update(int x,int v){
    int u=top[x];
    while (fa[u]){
        LL sum=Ts.query(w[u],w[u]+size[u]-1,1,n,1);
        Td.update(w[fa[u]],(sqr(val[x]-v)-sum*2*(val[x]-v) % mo)%mo,1,n,1);
        u=top[fa[u]];
    }
    Ts.update(w[x],v-val[x],1,n,1);
    val[x]=v;
}
LL query(int x,int y){
    LL res=0;
    while (top[x]!=top[y]){
        if (dep[top[x]]<dep[top[y]]) swap(x,y);
        res += Td.query(w[top[x]],w[x],1,n,1);
        res = res % mo;
       
        if (son[x]){
            LL sum=Ts.query(w[son[x]],w[son[x]]+size[son[x]]-1,1,n,1);
            res = res + sum*sum;
            res = res % mo;
        }
        LL sum=Ts.query(w[top[x]],w[top[x]]+size[top[x]]-1,1,n,1);
        
        res = res - sum*sum;
        res = res % mo;
        while (res<0) res+=mo;
        x=fa[top[x]];
    }
    if (dep[x]>dep[y]) swap(x,y);
    res += Td.query(w[x],w[y],1,n,1);
    res = res % mo;
    if (son[y]){
        LL sum=Ts.query(w[son[y]],w[son[y]]+size[son[y]]-1,1,n,1);
        res = res + sum*sum;
        res = res % mo;
    }
    if (fa[x]){
        LL sum=Ts.query(1,n,1,n,1)-Ts.query(w[x],w[x]+size[x]-1,1,n,1);
        res = res + sum*sum;
        res = res % mo;
    }
    return res;
}

int main(){
    while (~scanf("%d %d",&n,&m)){
        init();
        rep(i,1,n) scanf("%d",&val[i]);
        rep(i,2,n){
            int u,v;
            scanf("%d %d",&u,&v);
            add(u,v); add(v,u);
        }
        dfs_1(1);
        dfs_2(1,1);
        rep(i,1,n){
            int k=val[i];
            val[i]=0;
            update(i,k);
        }
        while (m--){
            int x,u,v;
            scanf("%d %d %d",&x,&u,&v);
            if (x==1) update(u,v);
                else {
                    LL sum=Ts.query(w[1],w[1]+size[1]-1,1,n,1);
                    sum = sum * sum;
                    sum = sum - query(u,v);
                    sum = sum % mo;
                    while (sum<0) sum+=mo;
                    printf("%lld\n",sum);
                }
        }
    }
}