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

http://www.cogs.pro/cogs/problem/problem.php?pid=2478

★☆   输入文件：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只得30，所以、、

发现，一个点u与连出去的点v，两者的LCA是u，所以在DFS过程中

sum记录当前点遍历次数，统计答案时用当前用所有的次数*当前点权

最后加上本身与本身的情况

#include <algorithm>
#include <cstdio>

#define LL long long

using namespace std;
 
const int N(1000010);
int n,w[N];
 
int sumedge,head[N];
struct Edge
{
    int from,to,next;
    Edge(int from=0,int to=0,int next=0):from(from),to(to),next(next){}
}edge[N<<1];
inline void ins(int from,int to)
{
    edge[++sumedge]=Edge(from,to,head[from]);
    head[from]=sumedge;
}

int dad[N];
LL ans,sum[N];
void DFS(int x)
{
    sum[x]=1;
    for(register int i=head[x];i;i=edge[i].next)
    {
        register int to=edge[i].to;
        if(to!=dad[x])
        {
            dad[to]=x; DFS(to);
            ans+=(LL)sum[x]*sum[to]*(LL)w[x];
            sum[x]+=(LL)sum[to];
        }
    }
}

inline void read(int &x)
{
    x=0;register char ch=getchar();
    for(;ch<'0'||ch>'9';) ch=getchar();
    for(;ch>='0'&&ch<='9';ch=getchar()) x=x*10+ch-'0';
}

int main()
{
    freopen("easy_LCA.in","r",stdin);
    freopen("easy_LCA.out","w",stdout);
    read(n);
    for(register int i=1;i<=n;i++) read(w[i]);
    for(register int u,v,i=1;i<n;i++)
        read(u),read(v),ins(u,v),ins(v,u);
    DFS(1);
    for(register int i=1;i<=n;i++) ans+=(LL)w[i];
    printf("%lld",ans);
    return 0;
}