Skip to content

14247: 【原4247】树上环形求和

题目

题目描述

author: fur 原OJ链接:https://acm.sjtu.edu.cn/OnlineJudge-old/problem/4247

树上环形求和

题目描述

有个一\(n\)个点\(n - 1\)条边的连通图,或者说一棵树,每个节点上有一个权值\(v_i\)。
现在想要知道,对于每个点,和它距离为\(2\)的所有点的权值之和

输入格式

第一行两个数\(n\),表示节点数
第二行\(n\)个数,分别表示每个节点的权值
接下来\(n - 1\)行,每行两个数\(a_i, b_i\),表示一条从\(a_i\)到\(b_i\)的连边

输出格式

一行\(n\)个数表示与节点\(0\)到\(n-1\)的距离为\(2\)的点的权值之和

样例输入1

5
1 2 3 4 5
0 1
1 2
2 3
3 4

样例输出1

3 4 6 2 3

数据规模

对于100%的数据有
\(\sum v_i \in range(int)\)
对于30%的数据有
\(n \leq 100\)
对于50%的数据有
\(n \leq 3000\)
对于另外20%的数据有
数据是随机的
对于100%的数据有
\(n \leq 100000\)

zqy2018's solution

#include <bits/stdc++.h>
#define INF 2000000000
using namespace std;
typedef long long ll;
int read(){
    int f = 1, x = 0;
    char c = getchar();
    while(c < '0' || c > '9'){if(c == '-') f = -f; c = getchar();}
    while(c >= '0' && c <= '9')x = x * 10 + c - '0', c = getchar();
    return f * x; 
}
int to[200005], nxt[200005], at[100005] = {0}, cnt = 0;
int n, a[100005], fa[100005] = {0};
ll ans[100005] = {0}, siz[100005] = {0};
void dfs(int cur, int faa){
    fa[cur] = faa;
    for (int i = at[cur]; i; i = nxt[i]){
        int v = to[i];
        if (v == faa) continue;
        siz[cur] += a[v], dfs(v, cur);
    }
}
void dfs2(int cur){
    if (fa[cur] != 0){
        ans[cur] += siz[fa[cur]] - a[cur];
        ans[cur] += a[fa[fa[cur]]];
    }
    for (int i = at[cur]; i; i = nxt[i]){
        int v = to[i];
        if (v == fa[cur]) continue;
        ans[cur] += siz[v], dfs2(v);
    }
}
void init(){
    n = read();
    for (int i = 1; i <= n; ++i)
        a[i] = read();
    for (int i = 1; i < n; ++i){
        int u = read() + 1, v = read() + 1;
        to[++cnt] = v, nxt[cnt] = at[u], at[u] = cnt;
        to[++cnt] = u, nxt[cnt] = at[v], at[v] = cnt;
    }
    a[0] = 0;
}
void solve(){
    dfs(1, 0);
    dfs2(1);
    for (int i = 1; i < n; ++i)
        printf("%lld ", ans[i]);
    printf("%lld\n", ans[n]);
}
int main(){
    init();
    solve();
    return 0;
}