Fedor Runs for President (Codeforces 1179D)
Tags: DP, trees

Fedor runs for president of Byteland! In the debates, he will be asked how to solve Byteland's transport problem. It's a really hard problem because of Byteland's transport system is now a tree (connected graph without cycles). Fedor's team has found out in the ministry of transport of Byteland that there is money in the budget only for one additional road. In the debates, he is going to say that he will build this road as a way to maximize the number of distinct simple paths in the country. A simple path is a path which goes through every vertex no more than once. Two simple paths are named distinct if sets of their edges are distinct. But Byteland's science is deteriorated, so Fedor's team hasn't succeeded to find any scientists to answer how many distinct simple paths they can achieve after adding exactly one edge on the transport system?Help Fedor to solve it.An edge can be added between vertices that are already connected, but it can't be a loop.In this problem, we consider only simple paths of length at least two.

Input: The first line contains one integer $$$n$$$ ($$$2 \leq n \leq 500\ 000$$$) — number of vertices in Byteland's transport system.Each of the following $$$n - 1$$$ lines contains two integers $$$v_i$$$ and $$$u_i$$$ ($$$1 \leq v_i, u_i \leq n$$$). It's guaranteed that the graph is tree.

Output: Print exactly one integer — a maximal number of simple paths that can be achieved after adding one edge.


Sample Inputs:
2
1 2

4
1 2
1 3
1 4

6
1 2
1 3
3 4
3 5
4 6

Sample Outputs:
2

11

29

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

def bfs(s):
    (q, k) = ([s], 0)
    visit = [0] * (n + 1)
    visit[s] = 1
    child = [[] for _ in range(n + 1)]
    while len(q) ^ k:
        i = q[k]
        for j in G[i]:
            if not visit[j]:
                q.append(j)
                visit[j] = 1
                child[i].append(j)
        k += 1
    return (q, child)
n = int(input())
G = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    (u, v) = map(int, input().split())
    G[u].append(v)
    G[v].append(u)
(p, child) = bfs(1)
p.reverse()
cnt = [1] * (n + 1)
for i in p:
    c = cnt[i]
    for j in child[i]:
        c += cnt[j]
    cnt[i] = c
x = [i * (i - 1) // 2 + 0.0 for i in range(n + 1)]
inf = pow(10, 12) + 1.0
dp = [0] * (n + 1)
xn = x[n]
mi = xn
for i in p:
    if not child[i]:
        continue
    dpi = inf
    c = cnt[i]
    for j in child[i]:
        mi = min(mi, dp[j] + x[n - cnt[j]])
        dpi = min(dpi, dp[j] + x[c - cnt[j]])
    dp[i] = dpi
    if len(child[i]) < 2:
        continue
    now = xn
    y = [xn - x[n - cnt[j]] - dp[j] for j in child[i]]
    j = y.index(max(y))
    now -= y[j]
    u = child[i][j]
    cu = n - cnt[u]
    y = [x[cu] - x[cu - cnt[j]] - dp[j] for j in child[i]]
    y[j] = -inf
    now -= max(y)
    mi = min(mi, now)
ans = int(2 * xn - mi)
print(ans)

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

def bfs(s):
    (q, k) = ([s], 0)
    visit = [0] * (n + 1)
    visit[s] = 1
    child = [[] for _ in range(n + 1)]
    while len(q) ^ k:
        i = q[k]
        for j in G[i]:
            if not visit[j]:
                q.append(j)
                visit[j] = 1
                child[i].append(j)
        k += 1
    return (q, child)
n = int(input())
G = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    (u, v) = map(int, input().split())
    G[u].append(v)
    G[v].append(u)
(p, child) = bfs(1)
p.reverse()
cnt = [1] * (n + 1)
for i in p:
    c = cnt[i]
    for j in child[i]:
        c += cnt[j]
    cnt[i] = c
x = [i * (i - 1) // 2 for i in range(n + 1)]
inf = pow(10, 12) + 1
dp = [0] * (n + 1)
mi = x[n]
for i in p:
    if not child[i]:
        continue
    dpi = inf
    c = cnt[i]
    for j in child[i]:
        mi = min(mi, dp[j] + x[n - cnt[j]])
        dpi = min(dpi, dp[j] + x[c - cnt[j]])
    dp[i] = dpi
    if len(child[i]) < 2:
        continue
    now = x[n]
    y = [x[n] - x[n - cnt[j]] - dp[j] for j in child[i]]
    j = y.index(max(y))
    now -= y[j]
    u = child[i][j]
    cu = n - cnt[u]
    y = [x[cu] - x[cu - cnt[j]] - dp[j] for j in child[i]]
    y[j] = -inf
    now -= max(y)
    mi = min(mi, now)
ans = 2 * x[n] - mi
print(ans)

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

def bfs(s):
    (q, k) = ([s], 0)
    visit = [0] * (n + 1)
    visit[s] = 1
    child = [[] for _ in range(n + 1)]
    while len(q) ^ k:
        i = q[k]
        for j in G[i]:
            if not visit[j]:
                q.append(j)
                visit[j] = 1
                child[i].append(j)
        k += 1
    return (q, child)
n = int(input())
G = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    (u, v) = map(int, input().split())
    G[u].append(v)
    G[v].append(u)
(p, child) = bfs(1)
p.reverse()
cnt = [1] * (n + 1)
for i in p:
    c = cnt[i]
    for j in child[i]:
        c += cnt[j]
    cnt[i] = c
x = [i * (i - 1) // 2 for i in range(n + 1)]
inf = pow(10, 12) + 1
dp = [0] * (n + 1)
mi = x[n]
for i in p:
    if not child[i]:
        continue
    dpi = inf
    c = cnt[i]
    for j in child[i]:
        mi = min(mi, dp[j] + x[n - cnt[j]])
        dpi = min(dpi, dp[j] + x[c - cnt[j]])
    dp[i] = dpi
    if len(child[i]) < 2:
        continue
    now = x[n]
    y = [x[n] - x[n - cnt[j]] - dp[j] for j in child[i]]
    j = y.index(max(y))
    now -= y[j]
    u = child[i][j]
    cu = n - cnt[u]
    y = [x[cu] - x[cu - cnt[j]] - dp[j] for j in child[i]]
    y[j] = -inf
    now -= max(y)
    mi = min(mi, now)
ans = 2 * x[n] - mi
print(ans)