1. 程式人生 > >[Luogu3676]小清新數據結構題

[Luogu3676]小清新數據結構題

www 等於 pro http next 現在 mat .com urn

題面戳我
題意:給一棵樹,樹上有點權,每次操作為修改一個點的點權,或者是詢問以某個點為根時,每棵子樹(以每個點為根,就有n棵子樹)點權和的平方和。
\(n\le2*10^5\),保證答案在long long範圍內

sol

我們設\(s_i\)表示以\(p\)為整棵樹的根時,以\(i\)為根的子樹的點權和。設\(Sum\)表示所有點的點權和,即\(Sum=\sum_{i=1}^{n}val_i\)
所以這道題給出\(p\),就是要你求\(\sum_{i=1}^{n}s_i^2\)

我們先看\(\sum_{i=1}^{n}s_i\)怎麽求。
考慮每個點的點權對\(\sum_{i=1}^{n}s_i\)

的貢獻,可以發現,每個點被計算了\(dep_i+1\)次,也就是說\(\sum_{i=1}^{n}s_i=\sum_{i=1}^{n}val_i(dep_i+1)=\sum_{i=1}^{n}val_idep_i+Sum\)。前面那一坨是不是有點熟悉?【ZJOI2015】幻想鄉戰略遊戲。
下文中為了方便描述,令\(calc(p)\)表示以\(p\)為根時的\(\sum_{i=1}^{n}val_idep_i\)

接下來我們考慮一下這個東西
\[\sum_{i=1}^{n}\sum_{j=1}^{n}val_ival_jdis(i,j)\]
這個可以形象地理解為,在每一對點對\((i,j)\)的路徑上每一條邊(剛好是\(dis(i,j)\)

條邊)上都加上\(val_ival_j\),然後求整棵樹上的邊權之和。

現在我們考慮每一條邊上的權值,它應該等於它兩側連接的兩坨樹的點權和的乘積。而連接的這兩坨樹中,不論取哪個\(p\)為根,都有有且僅有一坨樹會是一棵子樹。所以這個權值會等於\(s_i(Sum-s_i)\)。所以
\[\sum_{i=1}^{n}\sum_{j=1}^{n}val_ival_jdis(i,j)=\sum_{i=1}^{n}s_i(Sum-s_i)\]
這同時也證明了不論取哪個\(p\)作為根,\(\sum_{i=1}^{n}s_i(Sum-s_i)\)都不會變。

\(W=\sum_{i=1}^{n}s_i(Sum-s_i)\)

,可以先\(O(n)\)\(DP\)\(W\)的初值,然後就只要考慮一個點權修改對\(W\)的影響。
因為\(W=\sum_{i=1}^{n}\sum_{j=1}^{n}val_ival_jdis(i,j)\),若節點\(i\)的點權的變化量為\(\Delta v\),那麽\(\Delta W=\Delta v\sum_{j=1}^{n}val_jdis(i,j)\),相當於\(\Delta v*calc(i)\),所以說一樣地計算即可。
所以最終詢問的答案就是:
\[\sum_{i=1}^{n}s_i^2=Sum*\sum_{i=1}^{n}s_i-W=Sum(calc(i)+Sum)-W\]

code

#include<cstdio>
#include<algorithm>
using namespace std;
#define ll long long
const int N = 200005;
int gi()
{
    int x=0,w=1;char ch=getchar();
    while ((ch<'0'||ch>'9')&&ch!='-') ch=getchar();
    if (ch=='-') w=0,ch=getchar();
    while (ch>='0'&&ch<='9') x=(x<<3)+(x<<1)+ch-'0',ch=getchar();
    return w?x:-x;
}
struct edge{int to,next;}a[N<<1];
int n,q,val[N],head[N],cnt,pa[N],dep[N],sz[N],son[N],top[N];
void dfs1(int u,int f)
{
    pa[u]=f;dep[u]=dep[f]+1;sz[u]=1;
    for (int e=head[u];e;e=a[e].next)
    {
        int v=a[e].to;if (v==f) continue;
        dfs1(v,u);
        sz[u]+=sz[v];if (sz[v]>sz[son[u]]) son[u]=v;
    }
}
void dfs2(int u,int f)
{
    top[u]=f;
    if (son[u]) dfs2(son[u],f);else return;
    for (int e=head[u];e;e=a[e].next)
        if (a[e].to!=pa[u]&&a[e].to!=son[u])
            dfs2(a[e].to,a[e].to);
}
int lca(int u,int v)
{
    while (top[u]^top[v])
    {
        if (dep[top[u]]<dep[top[v]]) swap(u,v);
        u=pa[top[u]];
    }
    return dep[u]<dep[v]?u:v;
}
int getdis(int u,int v){return dep[u]+dep[v]-(dep[lca(u,v)]<<1);}
int tot,root,vis[N],w[N],fa[N];
ll sum[N],gather[N],tofa[N],sigma,omega,ans;
void getroot(int u,int f)
{
    sz[u]=1;w[u]=0;
    for (int e=head[u];e;e=a[e].next)
    {
        int v=a[e].to;if (v==f||vis[v]) continue;
        getroot(v,u);
        sz[u]+=sz[v];w[u]=max(w[u],sz[v]);
    }
    w[u]=max(w[u],tot-sz[u]);
    if (w[u]<w[root]) root=u;
}
void solve(int u,int f)
{
    fa[u]=f;vis[u]=1;
    for (int e=head[u];e;e=a[e].next)
    {
        int v=a[e].to;if (vis[v]) continue;
        tot=sz[v];
        root=0;
        getroot(v,0);
        solve(root,u);
    }
}
void modify(int u,int v)
{
    sum[u]+=v;
    for (int i=u;fa[i];i=fa[i])
    {
        int dist=getdis(u,fa[i]);
        sum[fa[i]]+=v;
        gather[fa[i]]+=dist*v;
        tofa[i]+=dist*v;
    }
}
ll calc(int u)
{
    ll res=gather[u];
    for (int i=u;fa[i];i=fa[i])
    {
        int dist=getdis(u,fa[i]);
        res+=(ll)dist*(sum[fa[i]]-sum[i]);
        res+=gather[fa[i]]-tofa[i];
    }
    return res;
}
void DP(int u)
{
    sz[u]=val[u];
    for (int e=head[u];e;e=a[e].next)
    {
        int v=a[e].to;if (v==pa[u]) continue;
        DP(v);sz[u]+=sz[v];
    }
    omega+=1ll*sz[u]*(sigma-sz[u]);
}
int main()
{
    n=gi();q=gi();
    for (int i=1;i<n;i++)
    {
        int u=gi(),v=gi();
        a[++cnt]=(edge){v,head[u]};head[u]=cnt;
        a[++cnt]=(edge){u,head[v]};head[v]=cnt;
    }
    dfs1(1,0);dfs2(1,1);
    tot=w[0]=n;
    getroot(1,0);
    solve(root,0);
    for (int i=1;i<=n;i++)
        val[i]=gi(),modify(i,val[i]),sigma+=val[i];
    DP(1);
    while (q--)
    {
        int opt=gi(),x=gi();
        if (opt==1)
        {
            int y=gi();
            modify(x,y-val[x]);sigma+=y-val[x];
            omega+=(y-val[x])*calc(x);
            val[x]=y;
        }
        else printf("%lld\n",(calc(x)+sigma)*sigma-omega);
    }
    return 0;
}

[Luogu3676]小清新數據結構題