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