Untitled

Anonymous
plain_text
02/28/2026 6:59 AM
2.0 KB
34
Indexable
import sys

def solve():
    sys.setrecursionlimit(300005)
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    n = int(input_data[0])
    k = int(input_data[1])
    
    val = [0] * (n + 1)
    for i in range(1, n + 1):
        val[i] = int(input_data[1 + i])
        
    adj = [[] for _ in range(n + 1)]
    idx = 2 + n
    for _ in range(n - 1):
        u = int(input_data[idx])
        v = int(input_data[idx+1])
        adj[u].append(v)
        adj[v].append(u)
        idx += 2
        
    dp1 = [1] * (n + 1)
    dp2 = [1] * (n + 1)
    ans = 1
    
    def dfs(u, p):
        nonlocal ans
        for v in adj[u]:
            if v != p:
                dfs(v, u)
                
        for v in adj[u]:
            if v != p:
                if val[v] - val[u] == k:
                    dp1[u] = max(dp1[u], 1 + dp1[v])
                if val[u] - val[v] == k:
                    dp2[u] = max(dp2[u], 1 + dp2[v])
                    
        if k == 0:
            m1, m2 = 0, 0
            for v in adj[u]:
                if v != p and val[v] == val[u]:
                    if dp1[v] > m1:
                        m2 = m1
                        m1 = dp1[v]
                    elif dp1[v] > m2:
                        m2 = dp1[v]
            ans = max(ans, m1 + m2 + 1)
        else:
            ans = max(ans, dp1[u] + dp2[u] - 1)
            
    dfs(1, 0)
    print(ans)

if __name__ == '__main__':
    solve()
Editor is loading...
Leave a Comment