Как рассчитать количество пар в дереве на основе максимального веса ребра? ⇐ Python

[Anonymous]

Я работаю с взвешенным деревом и мне нужно эффективно отвечать на несколько запросов. Каждый запрос запрашивает количество пар вершин (u,v) таких, что максимальный вес ребра на простом пути между u и v не превышает заданное значение 𝑞𝑖. Выходными данными должны быть результаты всех запросов по порядку.< /p>
Определения:
Дерево:

• Связный граф с n вершинами и циклов нет.
• Состоит из n – 1 ребер.
• Каждое ребро имеет вес 𝑤𝑖 и соединяет две вершины ui и vi.

Простой путь:

• Путь между двумя вершинами, который не повторно посетить любую вершину.​

Запрос

• По целому числу ци определите количество пар ( u,v), где u < v, так что максимальный вес ребра на простом пути между u и v равен qi.

Ввод:
1. Первая строка содержит два целых числа. n,m, где n — количество вершин, а m — количество запросов (1≤n,m≤200,000).
2. Каждая из следующих n-1 строк описывает ребро. в дереве с тремя целыми числами wi,ui,vi

• w𝑖: вес ребра (1≤𝑤𝑖≤200,000).
• ui,vi: вершины, соединенные ребром(1≤ui,vi≤n).

3.Последняя строка содержит m целые числа q1,q2,…,qm, где qi — максимальный вес или 𝑖-й запрос (1≤q𝑖≤200,000).

Вывод:
Одна строка, содержащая 𝑚 целые числа. Каждое целое число соответствует результату 𝑖-го запроса:

• Количество пар вершин (𝑢,𝑣), где 𝑢
  0 0
  Ввод:[/b]
  3 3
  1 2 1
  2 3 2
  1 3 2
  Выход:
  1 3 3
  
  Моя попытка:
  Я попробовал реализовать решение с помощью объединения-поиска (объединение непересекающихся множеств) с сортировкой ребер и обработкой запросов. Общий подход таков:
  • Сортировка ребер по весу.
  • Сортируйте запросы по 𝑞𝑖.
  • Используйте объединение-поиск для постепенного соединения узлов по весам ребер до 𝑞𝑖. >
  • Подсчитайте пары, используя размеры соединяемых компонентов.
  Вот мой код Python, который не подходит для тестовых случаев на моей университетской платформе:

  Код: Выделить всё

  from collections import defaultdict
  
  class UnionFind:
  def __init__(self, size):
  self.parent = list(range(size))
  self.size = [1] * size
  
  def find(self, x):
  if self.parent[x] != x:
  self.parent[x] = self.find(self.parent[x])
  return self.parent[x]
  
  def union(self, x, y):
  root_x = self.find(x)
  root_y = self.find(y)
  if root_x != root_y:
  if self.size[root_x] < self.size[root_y]:
  root_x, root_y = root_y, root_x
  self.parent[root_y] = root_x
  self.size[root_x] += self.size[root_y]
  return self.size[root_x] * self.size[root_y]
  return 0
  
  def count_pairs(n, m, edges, queries):
  edges.sort()
  queries = [(q, i) for i, q in enumerate(queries)]
  queries.sort()
  
  uf = UnionFind(n + 1)
  results = [0] * m
  current_pairs = 0
  edge_index = 0
  
  for q, idx in queries:
  while edge_index < len(edges) and edges[edge_index][0] 
  
  Подробнее здесь: [url]https://stackoverflow.com/questions/79305451/how-to-calculate-the-number-of-pairs-in-a-tree-based-on-maximum-edge-weight[/url]