Maximum White Subtree (Codeforces 1324F)
Tags: DP, trees

You are given a tree consisting of $$$n$$$ vertices. A tree is a connected undirected graph with $$$n-1$$$ edges. Each vertex $$$v$$$ of this tree has a color assigned to it ($$$a_v = 1$$$ if the vertex $$$v$$$ is white and $$$0$$$ if the vertex $$$v$$$ is black).You have to solve the following problem for each vertex $$$v$$$: what is the maximum difference between the number of white and the number of black vertices you can obtain if you choose some subtree of the given tree that contains the vertex $$$v$$$? The subtree of the tree is the connected subgraph of the given tree. More formally, if you choose the subtree that contains $$$cnt_w$$$ white vertices and $$$cnt_b$$$ black vertices, you have to maximize $$$cnt_w - cnt_b$$$.

Input: The first line of the input contains one integer $$$n$$$ ($$$2 \le n \le 2 \cdot 10^5$$$) — the number of vertices in the tree.The second line of the input contains $$$n$$$ integers $$$a_1, a_2, \dots, a_n$$$ ($$$0 \le a_i \le 1$$$), where $$$a_i$$$ is the color of the $$$i$$$-th vertex.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 $$$n$$$ integers $$$res_1, res_2, \dots, res_n$$$, where $$$res_i$$$ is the maximum possible difference between the number of white and black vertices in some subtree that contains the vertex $$$i$$$.

Note: The first example is shown below:The black vertices have bold borders.In the second example, the best subtree for vertices $$$2, 3$$$ and $$$4$$$ are vertices $$$2, 3$$$ and $$$4$$$ correspondingly. And the best subtree for the vertex $$$1$$$ is the subtree consisting of vertices $$$1$$$ and $$$3$$$.

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

4
0 0 1 0
1 2
1 3
1 4

Sample Outputs:
2 2 2 2 2 1 1 0 2

0 -1 1 -1

Solution (Python):
import gc
import heapq
import itertools
import math
from collections import Counter, deque, defaultdict
from sys import stdout
import time
from math import factorial, log, gcd
import sys
from decimal import Decimal
import threading
from heapq import *
from fractions import Fraction
import bisect

def S():
    return sys.stdin.readline().split()

def I():
    return [int(i) for i in sys.stdin.readline().split()]

def II():
    return int(sys.stdin.readline())

def IS():
    return sys.stdin.readline().replace('\n', '')

def main():
    n = II()
    a = I()
    tree = [[] for _ in range(n)]
    for _ in range(n - 1):
        (u, v) = I()
        tree[u - 1].append(v - 1)
        tree[v - 1].append(u - 1)
    order = []
    queue = deque([0])
    p = [-1] * n
    d = [-1 if i == 0 else 1 for i in a]
    while queue:
        v = queue.pop()
        order.append(v)
        for u in tree[v]:
            if p[v] != u:
                p[u] = v
                queue.appendleft(u)
    order = order[::-1]
    up = [False] * n
    for i in range(n - 1):
        if d[order[i]] > 0:
            up[order[i]] = True
            d[p[order[i]]] += d[order[i]]
    ans = [0] * n
    ans[0] = d[0]
    for i in range(n - 2, -1, -1):
        idx = order[i]
        if not up[idx]:
            ans[idx] = max(d[idx], ans[p[idx]] + d[idx])
        else:
            ans[idx] = max(ans[p[idx]], d[idx])
    print(*ans)
if __name__ == '__main__':
    main()

Solution (Python):
import gc
import heapq
import itertools
import math
from collections import Counter, deque, defaultdict
from sys import stdout
import time
from math import factorial, log, gcd
import sys
from decimal import Decimal
import threading
from heapq import *
from fractions import Fraction
import bisect

def S():
    return sys.stdin.readline().split()

def I():
    return [int(i) for i in sys.stdin.readline().split()]

def II():
    return int(sys.stdin.readline())

def IS():
    return sys.stdin.readline().replace('\n', '')

def main():
    n = II()
    a = I()
    tree = [[] for _ in range(n)]
    for _ in range(n - 1):
        (u, v) = I()
        tree[u - 1].append(v - 1)
        tree[v - 1].append(u - 1)
    order = []
    queue = deque([0])
    p = [-1] * n
    d = [-1 if i == 0 else 1 for i in a]
    while queue:
        v = queue.pop()
        order.append(v)
        for u in tree[v]:
            if p[v] != u:
                p[u] = v
                queue.appendleft(u)
    order = order[::-1]
    up = [False] * n
    for i in range(n - 1):
        if d[order[i]] > 0:
            up[order[i]] = True
            d[p[order[i]]] += d[order[i]]
    ans = [0] * n
    ans[0] = d[0]
    for i in range(n - 2, -1, -1):
        idx = order[i]
        if not up[idx]:
            ans[idx] = max(d[idx], ans[p[idx]] + d[idx])
        else:
            ans[idx] = max(ans[p[idx]], d[idx])
    print(*ans)
if __name__ == '__main__':
    main()

Solution (Python):
n = int(input())
a = list(map(int, input().split()))
G = [list() for _ in range(n)]
for _ in range(n - 1):
    (u, v) = map(int, input().split())
    G[u - 1].append(v - 1)
    G[v - 1].append(u - 1)
dp1 = [x * 2 - 1 for x in a] + [0]
p = [-1] * n
q = [0]
while q:
    v = q.pop()
    for x in G[v]:
        if x == p[v]:
            continue
        p[x] = v
        q.append(x)
dp2 = [0] * n
q = [0]
while q:
    v = q[-1]
    if dp2[v]:
        q.pop()
        if dp1[v] > 0:
            dp1[p[v]] += dp1[v]
    else:
        for x in G[v]:
            if x == p[v]:
                continue
            q.append(x)
        dp2[v] = 1
dp2 = [0] * n
q = [0]
while q:
    v = q.pop()
    for x in G[v]:
        if x == p[v]:
            continue
        c = dp1[v]
        if dp1[x] > 0:
            c -= dp1[x]
        if c > 0:
            dp1[x] += c
        q.append(x)
print(*dp1[:n])