cogs 2478. [HZOI 2016]简单的最近公共祖先

2478. [HZOI 2016]简单的最近公共祖先

★☆   输入文件：easy_LCA.in   输出文件：easy_LCA.out   简单对比
时间限制：2 s   内存限制：128 MB

【题目描述】

 

给定一棵有n个节点的有根树，根节点为1，每个节点有一个权值wi，求

即求所有无序节点对的LCA的权值之和。

树的节点编号为1~n，LCA表示两节点的最近公共祖先，即在它们的所有公共祖先中离根节点最远的节点。

【输入格式】

第一行一个整数n，表示节点数。

第二行n个正整数，表示每个点的权值。

以下n-1行每行两个整数x,y，表示树上有一条边连接节点x和节点y。

【输出格式】

一个整数，表示答案。

【样例输入】

3
1 2 3
1 2
1 3

【样例输出】

9

【数据范围与约定】

对于30%的数据，n<=1000。

对于60%的数据，n<=100000。

对于100%的数据，1<=n<=1000000，0<wi<=1000000。

【来源】

HZOI 2016

思路：lca。

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define MAXN 3000000
using namespace std;
int n,tot;
long long ans;
int w[MAXN],size[MAXN],dad[MAXN];
int to[MAXN],head[MAXN],net[MAXN];
void add(int u,int v){
    to[++tot]=v;net[tot]=head[u];head[u]=tot;
    to[++tot]=u;net[tot]=head[v];head[v]=tot;
}
void dfs(int now){
    size[now]=1;
    for(int i=head[now];i;i=net[i])
        if(dad[now]!=to[i]){
            dad[to[i]]=now;
            dfs(to[i]);
            ans+=1ll*size[now]*size[to[i]]*w[now];
            size[now]+=size[to[i]];
        }
}
int main(){
    freopen("easy_LCA.in","r",stdin);
    freopen("easy_LCA.out","w",stdout);
    scanf("%d",&n);
    for(int i=1;i<=n;i++){
        scanf("%d",&w[i]);
        ans+=w[i];
    }
    for(int i=1;i<n;i++){
        int x,y;
        scanf("%d%d",&x,&y);
        add(x,y);
    }
    dfs(1);
    printf("%lld",ans);
}