Tree with Maximum Cost (Codeforces 1092F)
Tags: DP, trees

You are given a tree consisting exactly of $$$n$$$ vertices. Tree is a connected undirected graph with $$$n-1$$$ edges. Each vertex $$$v$$$ of this tree has a value $$$a_v$$$ assigned to it.Let $$$dist(x, y)$$$ be the distance between the vertices $$$x$$$ and $$$y$$$. The distance between the vertices is the number of edges on the simple path between them.Let's define the cost of the tree as the following value: firstly, let's fix some vertex of the tree. Let it be $$$v$$$. Then the cost of the tree is $$$\sum\limits_{i = 1}^{n} dist(i, v) \cdot a_i$$$.Your task is to calculate the maximum possible cost of the tree if you can choose $$$v$$$ arbitrarily.

Input: The first line contains one integer $$$n$$$, the number of vertices in the tree ($$$1 \le n \le 2 \cdot 10^5$$$).The second line of the input contains $$$n$$$ integers $$$a_1, a_2, \dots, a_n$$$ ($$$1 \le a_i \le 2 \cdot 10^5$$$), where $$$a_i$$$ is the value of the vertex $$$i$$$.Each of the next $$$n - 1$$$ lines describes an edge of the tree. Edge $$$i$$$ is denoted by two integers $$$u_i$$$ and $$$v_i$$$, the labels of vertices it connects ($$$1 \le u_i, v_i \le n$$$, $$$u_i \ne v_i$$$).It is guaranteed that the given edges form a tree.

Output: Print one integer — the maximum possible cost of the tree if you can choose any vertex as $$$v$$$.

Note: Picture corresponding to the first example: You can choose the vertex $$$3$$$ as a root, then the answer will be $$$2 \cdot 9 + 1 \cdot 4 + 0 \cdot 1 + 3 \cdot 7 + 3 \cdot 10 + 4 \cdot 1 + 4 \cdot 6 + 4 \cdot 5 = 18 + 4 + 0 + 21 + 30 + 4 + 24 + 20 = 121$$$.In the second example tree consists only of one vertex so the answer is always $$$0$$$.

Sample Inputs:
8
9 4 1 7 10 1 6 5
1 2
2 3
1 4
1 5
5 6
5 7
5 8

1
1337

Sample Outputs:
121

0

Solution (Python):
import traceback
import sys
sys.setrecursionlimit(200010)
try:
    n = int(input())
    a = [0]
    a.extend(list(map(int, input().split())))
    D = [[] for i in range(n + 1)]
    for i in range(n - 1):
        (e1, e2) = map(int, input().split())
        D[e1].append(e2)
        D[e2].append(e1)
    visited = [False for i in range(n + 1)]
    cost = [a[i] for i in range(n + 1)]
    parent = [0 for i in range(n + 1)]
    val = 0

    def dfs(s, depth):
        global visited
        global cost
        global val
        global a
        global D
        stack = [(s, depth)]
        while stack:
            (s, depth) = stack[-1]
            if visited[s]:
                stack.pop()
                cost[parent[s]] += cost[s]
                continue
            else:
                visited[s] = True
                val += depth * a[s]
            for i in D[s]:
                if not visited[i]:
                    parent[i] = s
                    stack.append((i, depth + 1))
    dfs(1, 0)
    max_cost = val
    visited = [False for i in range(n + 1)]
    cost[0] = sum(a)

    def trav(s, some_val):
        global cost
        global visited
        global max_cost
        global D
        stack = [(s, some_val)]
        while stack:
            (s, some_val) = stack.pop()
            visited[s] = True
            if some_val > max_cost:
                max_cost = some_val
            for i in D[s]:
                if not visited[i]:
                    stack.append((i, some_val + (cost[0] - cost[i]) - cost[i]))
    trav(1, val)
    print(max_cost)
except Exception as ex:
    traceback.print_tb(ex.__traceback__)
    print(ex)

Solution (Python):
from sys import stdin, setrecursionlimit
n = int(stdin.readline())
ay = [int(x) for x in stdin.readline().split()]
graph = [set() for x in range(n)]
for e in range(n - 1):
    (a, b) = [int(x) - 1 for x in stdin.readline().split()]
    graph[a].add(b)
    graph[b].add(a)
dists = {}
child = {}
bruh = sum(ay)
path = [(0, 0)]
for (x, p) in path:
    for y in graph[x]:
        if y != p:
            path.append((y, x))
for (x, parent) in path[::-1]:
    total = 0
    children = ay[x]
    for v in graph[x]:
        if v != parent:
            (c, sm) = (child[v], dists[v])
            children += c
            total += sm + c
    dists[x] = total
    child[x] = children
(b2, b3) = (child[0], dists[0])
for (x, parent) in path:
    for v in graph[x]:
        if v != parent:
            dists[v] = dists[x] + bruh - child[v] * 2
print(max([dists[y] for y in range(n)]))

Solution (Python):
from collections import defaultdict
from sys import stdin, setrecursionlimit
import threading
input = stdin.buffer.readline

def work():
    res = 0
    ans = 0
    setrecursionlimit(1 << 18)
    n = int(input())
    a = list(map(int, input().split()))
    a.insert(0, 0)
    graph = defaultdict(list)
    dist = dict()
    for v in range(n - 1):
        (a1, b) = map(int, input().split())
        graph[a1].append(b)
        graph[b].append(a1)

    def dfs(v, par, h=0):
        nonlocal res
        res += h * a[v]
        dist[v] = a[v]
        for node in graph[v]:
            if node != par:
                dfs(node, v, h + 1)
                dist[v] += dist[node]

    def dfs2(v, par):
        nonlocal ans, res
        ans = max(ans, res)
        for node in graph[v]:
            if node != par:
                res -= dist[node]
                dist[v] -= dist[node]
                res += dist[v]
                dist[node] += dist[v]
                dfs2(node, v)
                dist[node] -= dist[v]
                res -= dist[v]
                dist[v] += dist[node]
                res += dist[node]
    dfs(1, -1)
    dfs2(1, -1)
    print(ans)
setrecursionlimit(300050)
threading.stack_size(100000000)
thread = threading.Thread(target=work)
thread.start()