Сума відстаней#10761

Сума відстаней (Розбір)

Використаємо систему неперетинних множин. Спочатку кожну вершину занесемо в окрему множину. Переберемо ребра в порядку спадання їхніх ваг, для кожного ребра $(u, v)$ запустимо функцію $U ni o n (u, v)$ Union(u, v). Вона об'єднає множини, в яких знаходяться вершини $u$ та $v$ .

Розглянемо поточне ребро $(u, v)$ дерева з вагою $w$ . Нехай воно об'єднує множини з розмірами $s i z e [u]$ size[u] та $s i z e [v]$ size[v]. Усі ребра в цих двох множинах мають вагу не меншу, ніж $w$ . Ребро $(u, v)$ входить у $s i z e [u] \cdot s i z e [v]$ size[u] \cdot size[v] різних шляхів. Оскільки вага $w$ мінімальна, то відстані цих шляхів дорівнюють $w$ . Тобто внесок цих шляхів у суму довжин усіх відстаней становить $w \cdot s i z e [u] \cdot s i z e [v]$ w \cdot size[u] \cdot size[v].

Приклад

Розглянемо перший тест.

Граф містить три шляхи:

$1 \to 2$ 1 → 2 (відстань $1$ );
$1 \to 3$ 1 → 3 (відстань $3$ );
$2 \to 1 \to 3$ 2 → 1 → 3 (відстань $1$ );

Сума відстаней між усіма парами вершин дорівнює $1 + 3 + 1 = 5$ 1 + 3 + 1 = 5.

Розглянемо другий тест.

Кожну вершину віднесемо до окремої множини.

Перебираємо ребра в порядку спадання ваги. Першим буде ребро $(1, 3)$ . Виконуємо $U ni o n (1, 3)$ Union(1, 3). Відстань $7$ має $s i z e [1] \cdot s i z e [3] = 1 \cdot 1 = 1$ size[1] \cdot size[3] = 1 \cdot 1 = 1 шлях (з $1$ в $3$ ).

Наступним буде ребро $(4, 5)$ . Виконуємо $U ni o n (4, 5)$ Union(4, 5). Відстань $5$ має $s i z e [4] \cdot s i z e [5] = 1 \cdot 1 = 1$ size[4] \cdot size[5] = 1 \cdot 1 = 1 шлях (з $4$ в $5$ ).

Наступним буде ребро $(2, 4)$ . Виконуємо $U ni o n (2, 4)$ Union(2, 4). Відстань $3$ має $s i z e [2] \cdot s i z e [4] = 1 \cdot 2 = 2$ size[2] \cdot size[4] = 1 \cdot 2 = 2 шляхи (з $2$ в $4$ і з $2$ в $5$ ).

Наступним буде ребро $(1, 4)$ . Виконуємо $U ni o n (1, 4)$ Union(1, 4). Відстань $2$ має $s i z e [1] \cdot s i z e [4] = 2 \cdot 3 = 6$ size[1] \cdot size[4] = 2 \cdot 3 = 6 шляхів $(2 \to 4 \to 1, 4 \to 1, 5 \to 4 \to 1, 2 \to 4 \to 1 \to 3, 4 \to 1 \to 3, 5 \to 4 \to 1 \to 3)$ (2 → 4 → 1, 4 → 1, 5 → 4 → 1, 2 → 4 → 1 → 3, 4 → 1 → 3, 5 → 4 → 1 → 3).

Шукана сума відстаней дорівнює $7 + 5 + 6 + 12 = 30$ 7 + 5 + 6 + 12 = 30.

Реалізація алгоритму

Оголосимо масиви, використовувані системою неперетинних множин.

struct Edge
{
  int u, v, cost;
} temp;
vector<Edge> e;
vector<int> parent, ssize;

Функція — компаратор cmpGreater сортує ребра за спаданням ваги.

int cmpGreater(Edge a, Edge b)
{
  return a.cost > b.cost;
}

Функція Repr повертає представника множини, в якій знаходиться вершина $v$ .

int Repr(int v)
{
  if (v == parent[v]) return v;
  return parent[v] = Repr(parent[v]);
}

Функція Union об'єднує множини з елементами $x$ та $y$ . Реалізуємо евристику з розмірами множин.

void Union(int x, int y)
{
  x = Repr(x); y = Repr(y);
  if (x == y) return;
  if (ssize[x] < ssize[y]) swap(x, y);
  parent[y] = x;
  ssize[x] += ssize[y];
}

C++

Основна частина програми. Читаємо кількість вершин $n$ в графі.

scanf("%d", &n);

Ініціалізуємо масиви.

parent.resize(n + 1);
ssize.resize(n + 1);
for (i = 0; i <= n; i++)
{
  parent[i] = i;
  ssize[i] = 1;
}

C++

Читаємо ребра дерева. Зберігаємо їх в масиві $e$ .

for (i = 0; i < n - 1; i++)
{
  scanf("%d %d %d", &temp.u, &temp.v, &temp.cost);
  e.push_back(temp);
}

Сортуємо ребра за спаданням ваг.

sort(e.begin(), e.end(), cmpGreater);

В змінній $r es$ підраховуємо суму відстаней між усіма парами вершин.

res = 0;

Перебираємо ребра. Для кожного ребра $(u, v)$ вагою $cos t$ запускаємо функцію $U ni o n (u, v)$ Union(u, v). В $r es$ додаємо вклад усіх шляхів з відстанню $cos t$ .

for (i = 0; i < e.size(); i++)
{
  res += 1LL * e[i].cost * ssize[Repr(e[i].u)] * ssize[Repr(e[i].v)];
  Union(e[i].u, e[i].v);
}

Виводимо відповідь.

printf("%lld\n", res);

Python реалізація

Оголосимо клас $E d g e$ , що описує ребро графа.

class Edge:
  def __init__(self, u, v, cost):
    self.u = u
    self.v = v
    self.cost = cost