Compressed Tree (Codeforces 1901E)
Tags: Greedy, trees

You are given a tree consisting of $$$n$$$ vertices. A number is written on each vertex; the number on vertex $$$i$$$ is equal to $$$a_i$$$.You can perform the following operation any number of times (possibly zero):   choose a vertex which has at most $$$1$$$ incident edge and remove this vertex from the tree. Note that you can delete all vertices.After all operations are done, you're compressing the tree. The compression process is done as follows. While there is a vertex having exactly $$$2$$$ incident edges in the tree, perform the following operation:   delete this vertex, connect its neighbors with an edge. It can be shown that if there are multiple ways to choose a vertex to delete during the compression process, the resulting tree is still the same.Your task is to calculate the maximum possible sum of numbers written on vertices after applying the aforementioned operation any number of times, and then compressing the tree.

Input: The first line contains a single integer $$$t$$$ ($$$1 \le t \le 10^4$$$) — the number of test cases.The first line of each test case contains a single integer $$$n$$$ ($$$2 \le n \le 5 \cdot 10^5$$$) — the number of vertices.The second line contains $$$n$$$ integers $$$a_1, a_2, \dots, a_n$$$ ($$$-10^9 \le a_i \le 10^9$$$).Each of the next $$$n - 1$$$ lines describes an edge of the tree. Edge $$$i$$$ is denoted by two integers $$$v_i$$$ and $$$u_i$$$, the labels of vertices it connects ($$$1 \le v_i, u_i \le n$$$, $$$v_i \ne u_i$$$). These edges form a tree. Additional constraint on the input: the sum of $$$n$$$ over all test cases doesn't exceed $$$5 \cdot 10^5$$$.

Output: For each test case, print a single integer — the maximum possible sum of numbers written on vertices after applying the aforementioned operation any number of times, and then compressing the tree.


Sample Inputs:
3
4
1 -2 2 1
1 2
3 2
2 4
2
-2 -5
2 1
7
-2 4 -2 3 3 2 -1
1 2
2 3
3 4
3 5
4 6
4 7

Sample Outputs:
3
0
9

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def make_graph(n, m):
    (x, s) = ([0] * (2 * m), [0] * (n + 3))
    for i in range(0, 2 * m, 2):
        (u, v) = map(int, input().split())
        s[u + 2] += 1
        s[v + 2] += 1
        (x[i], x[i + 1]) = (u, v)
    for i in range(3, n + 3):
        s[i] += s[i - 1]
    G = [0] * (2 * m)
    for i in range(2 * m):
        j = x[i] + 1
        G[s[j]] = x[i ^ 1]
        s[j] += 1
    return (G, s)

def bfs(s):
    q = [s]
    dist = [inf] * (n + 1)
    dist[s] = 0
    parent = [0] * (n + 1)
    for i in q:
        di = dist[i]
        for v in range(s0[i], s0[i + 1]):
            j = G[v]
            if dist[j] == inf:
                q.append(j)
                dist[j] = di + 1
                parent[j] = i
    return (q, dist, parent)

def f(u, v):
    return u << 2 ^ v
t = int(input())
ans = []
inf = pow(10, 15) + 1
for _ in range(t):
    n = int(input())
    a = [0] + list(map(int, input().split()))
    (G, s0) = make_graph(n, n - 1)
    if n == 2:
        ans0 = max(a[1], 0) + max(a[2], 0)
        ans.append(ans0)
        continue
    (p, dist, parent) = bfs(1)
    dp = [-inf] * (n + 1 << 2)
    ans0 = 0
    for i in p[::-1]:
        dp0 = [-inf] * 4
        dp0[0] = a[i]
        for v in range(s0[i], s0[i + 1]):
            j = G[v]
            if parent[i] == j:
                continue
            ma = -inf
            for k in range(4):
                ma = max(ma, dp[f(j, k)]) if k ^ 1 else max(ma, dp[f(j, k)] - a[j])
            for k in range(3, -1, -1):
                l = min(k + 1, 3)
                dp0[l] = max(dp0[l], dp0[k] + ma)
        for j in range(4):
            dp[f(i, j)] = dp0[j]
    ans0 = 0
    for i in range(1, n + 1):
        for j in range(4):
            ans0 = max(ans0, dp[f(i, j)]) if j ^ 2 else max(ans0, dp[f(i, j)] - a[i])
    ans.append(ans0)
sys.stdout.write('\n'.join(map(str, ans)))

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def make_graph(n, m):
    (x, s) = ([0] * (2 * m), [0] * (n + 3))
    for i in range(0, 2 * m, 2):
        (u, v) = map(int, input().split())
        s[u + 2] += 1
        s[v + 2] += 1
        (x[i], x[i + 1]) = (u, v)
    for i in range(3, n + 3):
        s[i] += s[i - 1]
    G = [0] * (2 * m)
    for i in range(2 * m):
        j = x[i] + 1
        G[s[j]] = x[i ^ 1]
        s[j] += 1
    return (G, s)

def bfs(s):
    q = [s]
    dist = [inf] * (n + 1)
    dist[s] = 0
    parent = [0] * (n + 1)
    for i in q:
        di = dist[i]
        for v in range(s0[i], s0[i + 1]):
            j = G[v]
            if dist[j] == inf:
                q.append(j)
                dist[j] = di + 1
                parent[j] = i
    return (q, dist, parent)

def f(u, v):
    return u << 2 ^ v
t = int(input())
ans = []
inf = pow(10, 15) + 1
for _ in range(t):
    n = int(input())
    a = [0] + list(map(int, input().split()))
    (G, s0) = make_graph(n, n - 1)
    if n == 2:
        ans0 = max(a[1], 0) + max(a[2], 0)
        ans.append(ans0)
        continue
    (p, dist, parent) = bfs(1)
    dp = [-inf] * (n + 1 << 2)
    ans0 = 0
    for i in p[::-1]:
        dp0 = [-inf] * 4
        dp0[0] = a[i]
        for v in range(s0[i], s0[i + 1]):
            j = G[v]
            if parent[i] == j:
                continue
            ma = -inf
            for k in range(4):
                ma = max(ma, dp[f(j, k)]) if k ^ 1 else max(ma, dp[f(j, k)] - a[j])
            for k in range(3, -1, -1):
                l = min(k + 1, 3)
                dp0[l] = max(dp0[l], dp0[k] + ma)
        for j in range(4):
            dp[f(i, j)] = dp0[j]
    ans0 = 0
    for i in range(1, n + 1):
        for j in range(4):
            ans0 = max(ans0, dp[f(i, j)]) if j ^ 2 else max(ans0, dp[f(i, j)] - a[i])
    ans.append(ans0)
sys.stdout.write('\n'.join(map(str, ans)))

Solution (Python):
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def make_graph(n, m):
    (x, s) = ([0] * (2 * m), [0] * (n + 3))
    for i in range(0, 2 * m, 2):
        (u, v) = map(int, input().split())
        s[u + 2] += 1
        s[v + 2] += 1
        (x[i], x[i + 1]) = (u, v)
    for i in range(3, n + 3):
        s[i] += s[i - 1]
    G = [0] * (2 * m)
    for i in range(2 * m):
        j = x[i] + 1
        G[s[j]] = x[i ^ 1]
        s[j] += 1
    return (G, s)

def bfs(s):
    q = [s]
    dist = [inf] * (n + 1)
    dist[s] = 0
    parent = [0] * (n + 1)
    for i in q:
        di = dist[i]
        for v in range(s0[i], s0[i + 1]):
            j = G[v]
            if dist[j] == inf:
                q.append(j)
                dist[j] = di + 1
                parent[j] = i
    return (q, dist, parent)

def f(u, v):
    return u << 2 ^ v
t = int(input())
ans = []
inf = pow(10, 15) + 1
for _ in range(t):
    n = int(input())
    a = [0] + list(map(int, input().split()))
    (G, s0) = make_graph(n, n - 1)
    if n == 2:
        ans0 = max(a[1], 0) + max(a[2], 0)
        ans.append(ans0)
        continue
    (p, dist, parent) = bfs(1)
    dp = [-inf] * (n + 1 << 2)
    ans0 = 0
    for i in p[::-1]:
        dp0 = [-inf] * 4
        dp0[0] = a[i]
        for v in range(s0[i], s0[i + 1]):
            j = G[v]
            if parent[i] == j:
                continue
            ma = -inf
            for k in range(4):
                ma = max(ma, dp[f(j, k)]) if k ^ 1 else max(ma, dp[f(j, k)] - a[j])
            for k in range(3, -1, -1):
                l = min(k + 1, 3)
                dp0[l] = max(dp0[l], dp0[k] + ma)
        for j in range(4):
            dp[f(i, j)] = dp0[j]
    ans0 = 0
    for i in range(1, n + 1):
        for j in range(4):
            ans0 = max(ans0, dp[f(i, j)]) if j ^ 2 else max(ans0, dp[f(i, j)] - a[i])
    ans.append(ans0)
sys.stdout.write('\n'.join(map(str, ans)))