jabajke

Untitled

May 25th, 2023
21
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 6.21 KB | None | 0 0
  1. BIG = 1e18
  2.  
  3.  
  4. class BidirectedGraph:
  5.  
  6. def __init__(self, values):
  7. self.n = len(values)
  8. self.values = values # В инит передаем массив значений вершин
  9. self.edges = [[] for i in range(self.n)]
  10.  
  11. def add_edge(self, u, v): # Когда добавляется ребро, создаем ребро из u IN v, v IN u
  12. # тк это двунаправленный граф
  13. u -= 1
  14. v -= 1
  15. self.edges[u].append(v)
  16. self.edges[v].append(u)
  17.  
  18. def find_max_path_dfs(self, v, parent):
  19. # Рекурсия будет считать ответ для поддерева вершины v
  20. # И этот ответ это будет максимум из ответа всех поддеревьев
  21. # И какого-то пути из одного поддерева через нашу вершину в другой
  22. # Будем использовать вспомогательный массив dp
  23. # В котором будет хранится наилучший путь, если будем идти от этой вершины
  24. # Как-нибудь вниз
  25. mx1, mx2, ans = -BIG, -BIG, -BIG
  26. self.dp[v] = self.values[v] # Изначально это вершина графа
  27. for child in self.edges[v]: # Итерируемся по детям
  28. if child == parent: # Смотри, чтоб наш ребенок не был предком, чтоб не зациклиться
  29. continue
  30. child_ans = self.find_max_path_dfs(child, v) # Вызываем рекурсию,
  31. # И находим ответ поддерева чайлда
  32. # Смотрим, что является максимальным из нашего
  33. ans = max(ans, child_ans)
  34. # Пытаемся улучшить ответ для нашего поддерева
  35. self.dp[v] = max(
  36. self.dp[v],
  37. max(0, self.dp[child]) + self.values[v]
  38. )
  39. if mx1 <= self.dp[child]: # Находим два самых максимальных dp[child]
  40. mx2 = mx1
  41. mx1 = self.dp[child]
  42. elif mx2 < self.dp[child]:
  43. mx2 = self.dp[child]
  44.  
  45. ans = max(ans, max(0, mx1) + max(0, mx2) + self.values[v])
  46. # ответ -
  47. return ans
  48.  
  49. def find_max_path(self):
  50. self.dp = [-BIG for i in range(self.n)]
  51. return self.find_max_path_dfs(0, -1)
  52.  
  53.  
  54. class DirectedGraph:
  55. def __init__(self, values):
  56. self.n = len(values)
  57. self.values = values
  58. self.edges = [[] for i in range(self.n)]
  59. self.counter = [0 for i in range(self.n)]
  60.  
  61. def add_edge(self, u, v):
  62. u -= 1
  63. # Здесь добавляем одно ребро и увеличиваем counter вершины на 1 в ту, которую идем
  64. # Чтобы потом найти root
  65. v -= 1
  66. self.edges[u].append(v)
  67. self.counter[v] += 1
  68.  
  69. def find_max_path_dfs(self, v):
  70. self.dp[v] = self.values[v] # Изначально ставим значение самой вершины
  71. for child in self.edges[v]: # Итерируемся по детям
  72. self.find_max_path_dfs(child) # Запускаем рекурсию от ребенка
  73. self.dp[v] = max(self.dp[v], # Обновляем наше значение dp[v]
  74. max(0, self.dp[child]) + self.values[v])
  75.  
  76. def find_root(self):
  77. for i in range(self.n): # Находим вершину, в которую не ведет ни одно ребро
  78. if self.counter[i] == 0:
  79. return i
  80.  
  81. def find_max_path(self): # Ответом будем максимальное значение в dp[v]
  82. self.dp = [-BIG for i in range(self.n)]
  83. root = self.find_root()
  84. self.find_max_path_dfs(root)
  85. return max(self.dp)
  86.  
  87.  
  88. def bidirected_graph_tests():
  89. graph = BidirectedGraph([1, 2, 3])
  90. graph.add_edge(1, 2)
  91. graph.add_edge(1, 3)
  92. assert graph.find_max_path() == 6
  93.  
  94. graph = BidirectedGraph([1])
  95. assert graph.find_max_path() == 1
  96.  
  97. graph = BidirectedGraph([-1, 2, 3])
  98. graph.add_edge(1, 2)
  99. graph.add_edge(1, 3)
  100. assert graph.find_max_path() == 4
  101.  
  102. graph = BidirectedGraph([-1, 2, 3])
  103. graph.add_edge(2, 1)
  104. graph.add_edge(2, 3)
  105. assert graph.find_max_path() == 5
  106.  
  107. graph = BidirectedGraph([1, 2, -3])
  108. graph.add_edge(1, 2)
  109. graph.add_edge(1, 3)
  110. assert graph.find_max_path() == 3
  111.  
  112. graph = BidirectedGraph([-100, 2, 3, 5, 6])
  113. graph.add_edge(1, 2)
  114. graph.add_edge(1, 3)
  115. graph.add_edge(2, 4)
  116. graph.add_edge(2, 5)
  117. assert graph.find_max_path() == 13
  118.  
  119. graph = BidirectedGraph([-100, 2, 20, 5, 6])
  120. graph.add_edge(1, 2)
  121. graph.add_edge(1, 3)
  122. graph.add_edge(2, 4)
  123. graph.add_edge(2, 5)
  124. assert graph.find_max_path() == 20
  125.  
  126. print("bidirected successful")
  127.  
  128.  
  129. def directed_graph_tests():
  130. graph = DirectedGraph([1, 2, 3])
  131. graph.add_edge(1, 2)
  132. graph.add_edge(1, 3)
  133. assert graph.find_max_path() == 4
  134.  
  135. graph = DirectedGraph([1])
  136. assert graph.find_max_path() == 1
  137.  
  138. graph = DirectedGraph([-1, 2, 3])
  139. graph.add_edge(1, 2)
  140. graph.add_edge(1, 3)
  141. assert graph.find_max_path() == 3
  142.  
  143. graph = DirectedGraph([-1, 2, 3])
  144. graph.add_edge(2, 1)
  145. graph.add_edge(2, 3)
  146. assert graph.find_max_path() == 5
  147.  
  148. graph = DirectedGraph([1, 2, -3])
  149. graph.add_edge(1, 2)
  150. graph.add_edge(1, 3)
  151. assert graph.find_max_path() == 3
  152.  
  153. graph = DirectedGraph([-100, 2, 3, 5, 6])
  154. graph.add_edge(1, 2)
  155. graph.add_edge(1, 3)
  156. graph.add_edge(2, 4)
  157. graph.add_edge(2, 5)
  158. assert graph.find_max_path() == 8
  159.  
  160. graph = DirectedGraph([-100, 2, 20, 5, 6])
  161. graph.add_edge(1, 2)
  162. graph.add_edge(1, 3)
  163. graph.add_edge(2, 4)
  164. graph.add_edge(2, 5)
  165. assert graph.find_max_path() == 20
  166.  
  167. print("directed successful")
  168.  
  169.  
  170. bidirected_graph_tests()
  171. directed_graph_tests()
Advertisement
Add Comment
Please, Sign In to add comment