Submission #64790120

Source Code Expand

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
import sys
sys.setrecursionlimit(10**6)
import pypyjit
pypyjit.set_param('max_unroll_recursion=-1')
def all_direction_tree_dp(n, edges):
    G = [[] for _ in range(n)]
    for u, v in edges:
        G[u].append(v)
        G[v].append(u)
    dp = [0] * n
    ans = [0] * n
    def dfs1(v, parent):
        for w in G[v]:
            if w == parent:
                continue
            dfs1(w, v)
            dp[v] = max(dp[v], dp[w] + 1)   
    def dfs2(v, parent, up):
        ans[v] = max(dp[v], up)
        children = [w for w in G[v] if w != parent]
        m = len(children)
 
 
הההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההההה
XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

import sys
sys.setrecursionlimit(10**6)
import pypyjit
pypyjit.set_param('max_unroll_recursion=-1')
def all_direction_tree_dp(n, edges):
    G = [[] for _ in range(n)]
    for u, v in edges:
        G[u].append(v)
        G[v].append(u)
    dp = [0] * n
    ans = [0] * n
    def dfs1(v, parent):
        for w in G[v]:
            if w == parent:
                continue
            dfs1(w, v)
            dp[v] = max(dp[v], dp[w] + 1)   
    def dfs2(v, parent, up):
        ans[v] = max(dp[v], up)
        children = [w for w in G[v] if w != parent]
        m = len(children)
        cand = [dp[w] + 2 for w in children]       
        prefix = [0] * (m + 1)
        suffix = [0] * (m + 1)
        for i in range(m):
            prefix[i+1] = max(prefix[i], cand[i])
        for i in range(m-1, -1, -1):
            suffix[i] = max(suffix[i+1], cand[i])
        
        for i, w in enumerate(children):
            siblings_max = max(prefix[i], suffix[i+1])
            new_up = max(up + 1, siblings_max)
            dfs2(w, v, new_up)
    
    dfs1(0, -1)
    dfs2(0, -1, 0)
    
    return ans

ans=0
n1=int(input())
g1=[]
for _ in range(n1-1):
    u,v=map(int,input().split())
    u-=1
    v-=1
    g1.append((u,v))
n2=int(input())
g2=[]
for _ in range(n2-1):
    u,v=map(int,input().split())
    u-=1
    v-=1
    g2.append((u,v))
ans=n1*n2
for c in all_direction_tree_dp(n1,g1):
    ans+=n2*c
for c in all_direction_tree_dp(n2,g2):
    ans+=n1*c
print(ans)

Submission Info

Judge Result