Tenzing and Tree (Codeforces 1842F)
Tags: Greedy, shortest paths

Tenzing has an undirected tree of $$$n$$$ vertices.Define the value of a tree with black and white vertices in the following way. The value of an edge is the absolute difference between the number of black nodes in the two components of the tree after deleting the edge. The value of the tree is the sum of values over all edges.For all $$$k$$$ such that $$$0 \leq k \leq n$$$, Tenzing wants to know the maximum value of the tree when $$$k$$$ vertices are painted black and $$$n-k$$$ vertices are painted white.

Input: The first line of the input contains a single integer $$$n$$$ ($$$1\leq n\leq 5000$$$) — the number of vertices.The following $$$n-1$$$ lines of the input contains $$$2$$$ integers $$$u_i$$$ and $$$v_i$$$ ($$$1 \leq u_i, v_i \leq n, u_i \neq v_i$$$) — indicating an edge between vertices $$$u_i$$$ and $$$v_i$$$. It is guaranteed that the given edges form a tree.

Output: Output $$$n+1$$$ numbers. The $$$i$$$-th number is the answer of $$$k=i-1$$$.

Note: Consider the first example. When $$$k=2$$$, Tenzing can paint vertices $$$1$$$ and $$$2$$$ black then the value of edge $$$(1,2)$$$ is 0, and the values of other edges are all equal to $$$2$$$. So the value of that tree is $$$4$$$.

Sample Inputs:
4
1 2
3 2
2 4

1

Sample Outputs:
0 3 4 5 6

0 0

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
    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
    return (q, dist)
n = int(input())
(G, s0) = make_graph(n, n - 1)
inf = pow(10, 9) + 1
ans = [0] * (n + 1)
for s in range(1, n + 1):
    if n >= 3 and s0[s + 1] - s0[s] == 1:
        continue
    (p, dist) = bfs(s)
    u = 0
    for i in range(1, n + 1):
        u += dist[p[i - 1]]
        ans[i] = max(ans[i], (n - 1) * i - 2 * u)
sys.stdout.write(' '.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)
n = int(input())
(G, s0) = make_graph(n, n - 1)
inf = pow(10, 9) + 1
mi = [inf] * (n + 1)
q = [0] * n
for s in range(1, n + 1):
    if n >= 3 and s0[s + 1] - s0[s] == 1:
        continue
    (q[0], l, u) = (s, 1, 0)
    dist = [inf] * (n + 1)
    dist[s] = 0
    for i in q:
        di = dist[i]
        for v in range(s0[i], s0[i + 1]):
            j = G[v]
            if dist[j] == inf:
                q[l] = j
                dist[j] = di + 1
                u += di + 1
                mi[l + 1] = min(mi[l + 1], u)
                l += 1
ans = [(n - 1) * i - 2 * mi[i] for i in range(n + 1)]
(ans[0], ans[1]) = (0, n - 1)
sys.stdout.write(' '.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
    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
    return (q, dist)
n = int(input())
(G, s0) = make_graph(n, n - 1)
inf = pow(10, 9) + 1
ans = [0] * (n + 1)
for s in range(1, n + 1):
    (p, dist) = bfs(s)
    u = 0
    for i in range(1, n + 1):
        u += dist[p[i - 1]]
        ans[i] = max(ans[i], (n - 1) * i - 2 * u)
sys.stdout.write(' '.join(map(str, ans)))