8.5-Day1T2--Asm.Def 的基本算法

题目大意

给一棵树，求∑∑w_i*w_j*w_LCA(i,j)

w_i表示i点权值

题解

显然一点点求lca是肯定会tle的

那就想如何优化

i和j的lca和j和i的lca是一样的

DFS，在每个x处，统计以它为LCA的答案总和

假设x有k个子树，权值和分别是S1,S2,…,Sk

设P=S1+S2+...+Sk 这些值可以轻易地在DFS中求出

①i、j分别在两棵不同的子树：{S1(P-S1) + S2(P-S2) + … +Sk(P-Sk)}*W_x

②i、j之一是x：(2W_x*P)*W_x

③i、j都是x：(W_x^2)*W_x

#include<cstdio>
#include<algorithm>
#define ll long long
using namespace std;
inline ll read()
{
    ll sum = 0,p = 1;
    char ch = getchar();
    while(ch < '0' || ch > '9')
    {
        if(ch == '-')
            p = -1;
        ch = getchar();
    }
    while(ch >= '0' && ch <= '9')
    {
        (sum *= 10) += ch - '0';
        ch = getchar();
    }
    return sum * p;
}
const int N = 200005;
const ll mod = 1000000007;
int n;
ll w[N],sum[N],ans;
int head[N],cnt;
struct edge
{
    int to,nxt;
}e[N];
void add(int a,int b)
{
    e[++cnt].nxt= head[a];
    e[cnt].to = b;
    head[a] = cnt;
}
void dfs(int u,int fa)
{
    sum[u] = w[u];
    for(int i = head[u];i;i = e[i].nxt)
    {
        int v = e[i].to;
        if(v == fa)
            continue;
        dfs(v,u);
        (sum[u] += sum[v])%= mod;
    }
}
void lcaa(int u,int fa)
{
    for(int i = head[u];i;i = e[i].nxt)
    {
        int v = e[i].to;
        if(v == fa)
            continue;
        ans = (ans + w[u]%mod * w[u]%mod * sum[v]%mod)%mod;
        ans = (ans + (sum[u] - w[u] - sum[v])%mod *w[u]%mod * sum[v]%mod)%mod;
        sum[u] = (sum[u] - sum[v] + mod)%mod;
    }
    for(int i = head[u];i;i = e[i].nxt)
    {
        int v = e[i].to;
        if(v == fa)
            continue;
        lcaa(v,u);
    }
}
int main()
{
    n = read();
    w[1] = read();
    for(int i = 2;i <= n;i++)
    {
        int fa = read();
        w[i] = read();
        add(fa,i);
        add(i,fa);
    }
    dfs(1,0);
    lcaa(1,0);
    ans = (ans * (ll)2) % mod;
    for(int i = 1;i <= n;i++)
        ans = (ans + w[i]%mod * w[i]% mod * w[i]%mod)%mod;
    printf("%lld",ans);
    return 0;
}