fork download
copy 
  1. import sys
  2. input = [int(x) for x in sys.stdin.read().split()][::-1].pop
  3.  
  4.  
  5. def solve():
  6.     n = input()
  7.     s = input()
  8.     graph = [{} for _ in range(n)]
  9.     for _ in range(n-1):
  10.         x = input()
  11.         y = input()
  12.         z = input()
  13.         graph[x-1][y-1] = z
  14.         graph[y-1][x-1] = z
  15.  
  16.     # Do DFS to find the farthest node from 0
  17.     stack: list[tuple[int, int]] = [(0, 0)]
  18.     visited: list[bool] = [False] * n
  19.     visited[0] = True
  20.     max_dist: int = 0
  21.     max_node: int = 0
  22.     while stack:
  23.         node, dist = stack.pop()
  24.         if dist > max_dist:
  25.             max_dist = dist
  26.             max_node = node
  27.         for neighbor, weight in graph[node].items():
  28.             if not visited[neighbor]:
  29.                 visited[neighbor] = True
  30.                 stack.append((neighbor, dist + weight))
  31.  
  32.     # Do DFS to find the farthest node from max_node and get the path between them
  33.     stack = [(max_node, 0)]
  34.     visited = [False] * n
  35.     visited[max_node] = True
  36.     parent: list[int] = [-1] * n
  37.     max_dist = 0
  38.     second_max_node = 0
  39.     while stack:
  40.         node, dist = stack.pop()
  41.         if dist > max_dist:
  42.             max_dist = dist
  43.             second_max_node = node
  44.         for neighbor, weight in graph[node].items():
  45.             if not visited[neighbor]:
  46.                 visited[neighbor] = True
  47.                 stack.append((neighbor, dist + weight))
  48.                 parent[neighbor] = node
  49.  
  50.     longest_path: list[int] = []
  51.     curr = second_max_node
  52.     while curr != -1:
  53.         longest_path.append(curr)
  54.         curr = parent[curr]
  55.  
  56.     l: int = 0
  57.     r: int = len(longest_path) - 1
  58.     sum_left: int = 0
  59.     sum_right: int = 0
  60.     while l < r:
  61.         if sum_left < sum_right:
  62.             sum_left += graph[longest_path[l]][longest_path[l+1]]
  63.             l += 1
  64.         else:
  65.             sum_right += graph[longest_path[r]][longest_path[r-1]]
  66.             r -= 1
  67.  
  68.     while l > 0 or r < len(longest_path) - 1:
  69.         if sum_left == 0 and sum_right == 0:
  70.             break
  71.         if sum_left < sum_right:
  72.             # nullify the right side
  73.             if graph[longest_path[r]][longest_path[r+1]] <= s:
  74.                 s -= graph[longest_path[r]][longest_path[r+1]]
  75.                 sum_right -= graph[longest_path[r]][longest_path[r+1]]
  76.                 r += 1
  77.             else:
  78.                 break
  79.         else:
  80.             # nullify the left side
  81.             if graph[longest_path[l]][longest_path[l-1]] <= s:
  82.                 s -= graph[longest_path[l]][longest_path[l-1]]
  83.                 sum_left -= graph[longest_path[l]][longest_path[l-1]]
  84.                 l -= 1
  85.             else:
  86.                 break
  87.  
  88.     ans: int = max(sum_left, sum_right)
  89.  
  90.     # Do DFS from path to get longest distance
  91.     visited = [False] * n
  92.     for node in longest_path:
  93.         visited[node] = True
  94.     stack = [(node, 0) for node in longest_path]
  95.     while stack:
  96.         node, dist = stack.pop()
  97.         if dist > ans:
  98.             ans = dist
  99.         for neighbor, weight in graph[node].items():
  100.             if not visited[neighbor]:
  101.                 visited[neighbor] = True
  102.                 stack.append((neighbor, dist + weight))
  103.  
  104.     print(ans)
  105.  
  106.  
  107. if __name__ == "__main__":
  108.     n_tests: int = input()
  109.     for test_nb in range(n_tests):
  110.         solve()
  111.  
Success #stdin #stdout 0.03s 9784KB
comments (?)
stdin
copy 
2
5 2
1 2 5
2 3 2
2 4 4
2 5 3
8 6
1 3 2
2 3 2
3 4 6
4 5 3
4 6 4
4 7 2
7 8 3
stdout
copy 
5
5
https://ideone.com/f4jLVc
problem MINDIST, based on SPOJ submission 33476689 (PYPY3), 2024-09-11 18:01:31
language:
Python 3 (python  3.12)
created:
1 year ago
visibility:
secret

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!