Cut the Tree

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)