Atul 对图论感兴趣,最近他正在学习树形结构。他发现在一颗树 T 上移除一条边会生成两个分开的树:T1 和 T2。
树 T 的每个顶点被赋予了一个正整数。你的任务是移除一条边使得 *Tree_diff* 最小化, *Tree_diff* 的定义如下:
F(T) = 树T上各顶点的数值之和。
Tree_diff(T) = abs(F(T1) - F(T2))
输入格式
第一行包含一个整数 N,代表树的顶点个数。
第二行包含 N 个以空格隔开的整数,第i个数代表标号为i的顶点上的数值。
接下来的(N - 1)行每行代表一条边。
输出格式
用一行输出最小的 *Tree_diff*。
数据范围
3 ≤ N ≤ 105
1 ≤ 每个顶点上的数值 ≤ 1001
顶点标号都是从 1 到 N之间的整数。
样例输入
6
100 200 100 500 100 600
1 2
2 3
2 5
4 5
5 6
样例输出
400
*样例解释
切开 1 2
这条边, Tree_diff = 1400
切开 2 3
这条边, Tree_diff = 1400
切开 2 5
这条边, Tree_diff = 800
切开 4 5
这条边, Tree_diff = 600
切开 5 6
这条边, Tree_diff = 400
所以答案为 400.
思路:首先题目说的有问题,输入的是无向图,而非tree。但是思路是先转换成tree,然后求每个节点的所有子节点的和,BFS就好了
from collections import defaultdict
# Complete the solve function below.
def solve(n, data, edges):
# undirected graph --> tree
degree=[0]*(1+n)
adj=defaultdict(set)
for p,s in edges:
adj[s].add(p)
adj[p].add(s)
root=1
q,qq=[root],[]
vis=[False]*(1+n)
vis[root]=True
c2p=[-1]*(n+1)
while q:
while q:
p=q.pop()
for s in adj[p]:
if vis[s]: continue
vis[s]=True
qq.append(s)
c2p[s]=p
degree[p]+=1
q,qq=qq,q
q=[i for i in range(1,n+1) if degree[i]==0]
vis=[False]*(1+n)
res=[0]*(1+n)
for i in q: vis[i]=True
for i in range(1,n+1): res[i]=data[i-1]
while q:
s=q.pop()
p = c2p[s]
if vis[p]: continue
degree[p]-=1
res[p]+=res[s]
if degree[p]==0:
q.append(p)
vis[p]=True
# print(res)
su=max(res)
mi=res[1]
for i in range(2,n+1):
if abs(res[i]*2-su)<abs(mi*2-su):
mi=res[i]
return abs(mi*2-su)
if __name__ == '__main__':
n = int(input())
data = list(map(int, input().rstrip().split()))
edges = []
for _ in range(n-1):
edges.append(list(map(int, input().rstrip().split())))
result = solve(n, data, edges)
print(result)