Məsafələrin cəmi (Həll təhlili)

Kesişməyən çoxluqlar sistemindən istifadə edəcəyik. Əvvəlcə hər zirvəni ayrı bir çoxluğa daxil edəcəyik. Kənarları onların çəkilərinin azalan sırası ilə yoxlayacağıq, hər kənar üçün $(u,v)$ funksiyasını işə salacağıq. Bu funksiyalar $u$ və $v$ zirvələrinin bulunduğu çoxluqları birləşdirəcək.

Ağacın çəkisi $w$ olan cərəyan $(u,v)$ kənarını nəzərdən keçirək. Təsəvvür edək ki, o, size[u] və size[v] ölçüləri olan çoxluqları birləşdirir. Bu iki çoxluqdan olan bütün kənarların çəkisi az olmayan $w$-dir. $(u,v)$ kənarı size[u] \cdot size[v] müxtəlif yollarda iştirak edir. Çəkisi $w$ ən kiçik olduğu üçün bu yolların məsafələri də $w$-yə bərabərdir. Başqa sözlə, bu yolların bütünlərdəki töhfəsi yolların uzunluqlarının cəmi w \cdot size[u] \cdot size[v]-dir.

Misal

İlk testi götürək.

Qrafda üç yol var:

• 1 → 2 (məsafə $1$);

• 1 → 3 (məsafə $3$);

• 2 → 1 → 3 (məsafə $1$);

Bütün zirvə cütləri arasında məsafələrin cəmi 1 + 3 + 1 = 5-ə bərabərdir.

İkinci testi götürək.

Hər zirvəni fərdi bir çoxluğa aid edək.

Kənarları azalan çəki sırası ilə yoxlayacağıq. İlk kənar $(1,3)$ olacaq. Union(1, 3)-ü icra edirik. $7$ mükafatlı məsafənin size[1] \cdot size[3] = 1 \cdot 1 = 1 yolu var (1-dən 3-ə).

Növbəti kənar $(4,5)$ olacaq. Union(4, 5)-ü icra edirik. $5$ mükafatlı məsafənin size[4] \cdot size[5] = 1 \cdot 1 = 1 yolu var (4-dən 5-ə).

Növbəti kənar $(2,4)$ olacaq. Union(2, 4)-ü icra edirik. $3$ mükafatlı məsafənin size[2] \cdot size[4] = 1 \cdot 2 = 2 yolu var (2-dən 4-ə və 2-dən 5-ə).

Sonrakı kənar $(1,4)$ olacaq. Union(1, 4)-ü icra edirik. $2$ mükafatlı məsafənin size[1] \cdot size[4] = 2 \cdot 3 = 6 yolu var (2 → 4 → 1, 4 → 1, 5 → 4 → 1, 2 → 4 → 1 → 3, 4 → 1 → 3, 5 → 4 → 1 → 3).

Axtarılan məsafələr cəmi 7 + 5 + 6 + 12 = 30-a bərabərdir.

Algoritmin tətbiqi

Kesişməyən çoxluqlar sistemi tərəfindən istifadə olunan massivləri elan edirik.

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

cmpGreater komparator funksiyası kənarları azalan sıra ilə çəkilərinə görə sıralayır.

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

Repr funksiyası $v$ zirvəsinin aid olduğu çoxluğun nümayəndəsini qaytarır.

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

Union funksiyası, $x$ və $y$ elementləri olan çoxluqları birləşdirir. Çoxluqların ölçüsü ilə bağlı evristika həyata keçirəcəyik.

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];
}

Proqramın əsas hissəsi. Qrafdakı zirvələrin sayını $n$ oxuyuruq.

scanf("%d", &n);

Massivləri başladırıq.

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

Ağacın kənarlarını oxuyuruq. Onları $e$ massivində saxlayırıq.

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

Kənarları azalan çəki sırası ilə sıralayırıq.

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

$res$ dəyişənində bütün zirvə cütləri arasındakı məsafələrin cəmini hesablayırıq.

res = 0;

Kənarları yoxlayırıq. Hər $(u,v)$ kənarı üçün $cost$ ağırlıqla Union(u, v) funksiyasını işlədirik. Bütün $cost$ ağırlıqlı yolların töhfəsini $res$-ə əlavə edirik.

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);
}

Nəticəni çap edirik.

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

Python tətbiqi

Qrafın kənarını təsvir edən $Edge$ klası elan edirik.

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

cmpGreater komparator funksiyası kənarları azalan sıra ilə çəkilərinə görə sıralayır.

def cmpGreater(a, b):
    return a.cost > b.cost

Repr funksiyası $v$ zirvəsinin aid olduğu çoxluğun nümayəndəsini qaytarır.

def Repr(v):
  if v == parent[v]: return v
  parent[v] = Repr(parent[v])
  return parent[v]

Union funksiyası, $x$ və $y$ elementləri olan çoxluqları birləşdirir. Çoxluqların ölçüsü ilə bağlı evristika həyata keçirəcəyik.

def Union(x, y):
  x = Repr(x)
  y = Repr(y)
  if x == y: return
  if ssize[x] < ssize[y]:
    x, y = y, x
  parent[y] = x
  ssize[x] += ssize[y]

Proqramın əsas hissəsi. Qrafdakı zirvələrin sayını $n$ oxuyuruq.

n = int(input().strip())

Massivləri başladırıq.

parent = list(range(n + 1))
ssize = [1] * (n + 1)
e = []

Ağacın kənarlarını oxuyuruq. Onları $e$ siyahısında saxlayırıq.

for _ in range(n - 1):
  u, v, cost = map(int, input().split())
  e.append(Edge(u, v, cost))

Kənarları azalan çəki sırası ilə sıralayırıq.

e.sort(key=lambda x: x.cost, reverse=True)

$res$ dəyişənində bütün zirvə cütləri arasındakı məsafələrin cəmini hesablayırıq.

res = 0

Kənarları yoxlayırıq. Hər $(u,v)$ kənarı üçün $cost$ ağırlıqla Union(u, v) funksiyasını işlədirik. Bütün $cost$ ağırlıqlı yolların töhfəsini $res$-ə əlavə edirik.

for edge in e:
  res += edge.cost * ssize[Repr(edge.u)] * ssize[Repr(edge.v)]
  Union(edge.u, edge.v)

Nəticəni çap edirik.

print(res)