Сума#2157

Сума (Розбір)

Запустимо пошук в глибину з довільної вершини дерева. Розглянемо ребро $(u, v)$ з вагою $w$ . Нехай кількість вершин у піддереві з коренем $v$ дорівнює $c$ . Тоді з одного боку ребра знаходиться $c$ вершин, а з іншого $n - c$ n - c вершин. Ребро $(u, v)$ входить в $c \cdot (n - c)$ c \cdot (n - c) різних шляхів між парами вершин. Відповідно, внесок ребра $(u, v)$ в сумарну довжину всіх шляхів дорівнює $c \cdot (n - c) \cdot w$ c \cdot (n - c) \cdot w. Залишається перебрати всі ребра за допомогою пошуку в глибину та підсумувати їх внесок у загальну суму довжин шляхів.

Приклад

Графи, наведені у прикладах, мають вигляд:

Розглянемо внесок ребер у загальну суму довжин всіх шляхів у першому прикладі.

Ребро $(1, 2)$ з вагою $1$ належить двом шляхам:

$1 - 2$ 1 - 2;
$2 - 1 - 3$ 2 - 1 - 3;

Його внесок у загальну суму дорівнює $1 \cdot 2 \cdot 1 = 2$ 1 \cdot 2 \cdot 1 = 2.

Ребро $(1, 3)$ з вагою $3$ належить двом шляхам:

$1 - 3$ 1 - 3;
$2 - 1 - 3$ 2 - 1 - 3;

Його внесок у загальну суму дорівнює $2 \cdot 1 \cdot 3 = 6$ 2 \cdot 1 \cdot 3 = 6.

Сумарна довжина всіх шляхів дорівнює $2 + 6 = 8$ 2 + 6 = 8.

У другому прикладі розглянемо внесок ребра $(1, 2)$ з вагою $5$ у загальну суму довжин усіх шляхів.

Кількість вершин у піддереві з коренем у вершині $2$ дорівнює $c = 4$ c = 4 (включаючи саму вершину $2$ ). Тоді з іншого боку ребра $(1, 2)$ знаходиться $n - c = 6 - 4 = 2$ n - c = 6 - 4 = 2 вершини. Відповідно, кількість різних шляхів, що проходять через ребро $(1, 2)$ , дорівнює $c * (n - c) = 4 * 2 = 8$ c * (n - c) = 4 * 2 = 8. Дійсно, такими шляхами є всі пари вершин, в яких одна вершина лежить у піддереві вершини $2$ , а інша — поза ним:

1 - 2, 1 - 2 - 4, 1 - 2 - 5, 1 - 2 - 6, 3 - 1 - 2, 3 - 1 - 2 - 4, 3 - 1 - 2 - 5, 3 - 1 - 2 - 6

1 - 2, 1 - 2 - 4, 1 - 2 - 5, 1 - 2 - 6, \\ 3 - 1 - 2, 3 - 1 - 2 - 4, 3 - 1 - 2 - 5, 3 - 1 - 2 - 6

Внесок ребра $(1, 2)$ з вагою $5$ у суму довжин всіх шляхів становить

c \cdot (n - c) \cdot w = 4 \cdot 2 \cdot 5 = 40

c \cdot (n - c) \cdot w = 4 \cdot 2 \cdot 5 = 40

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

Оголосимо значення модуля $M O D$ .

#define MOD 1000000000

Вхідний зважений граф зберігаємо у списку суміжності $g$ .

vector<vector<pair<int,int> > > g;

Функція dfs реалізує пошук в глибину з вершини $v$ і повертає кількість вершин у піддереві з коренем $v$ (включаючи саму вершину $v$ ). Підрахунок цих вершин ведеться у змінній $c n t$ . Ізначально вважаємо $c n t = 1$ cnt = 1, оскільки саме піддерево вже містить вершину $v$ .

int dfs(int v, int p = -1)
{
  int cnt = 1;
  for (auto edge : g[v])
  {
    int to = edge.first;
    int w = edge.second;

C++

Розглянемо ребро $(v, t o)$ (v, to) з вагою $w$ . Визначаємо кількість вершин $c$ у піддереві з коренем у вершині $t o$ . Тоді з одного боку ребра знаходяться $c$ вершин, а з іншого $n - c$ n - c вершин. Ребро $(v, t o)$ (v, to) входить в $c \cdot (n - c)$ c \cdot (n - c) різних шляхів між парами вершин.

Таким чином, внесок ребра $(v, t o)$ (v, to) у сумарну довжину всіх шляхів дорівнює

c \cdot (n - c) \cdot w

c \cdot (n - c) \cdot w

    if (to != p)
    {
      int c = dfs(to, v);
      res = (res + 1LL * c * (n - c) * w) % MOD;
      cnt += c;
    }
  }
  return cnt;
}

C++

Основна частина програми. Читаємо вхідний зважений граф у список суміжності $g$ .

scanf("%d", &n);
g.resize(n + 1);
for (i = 1; i < n; i++)
{
  scanf("%d %d %d", &u, &v, &d);
  g[u].push_back({v, d});
  g[v].push_back({u, d});
}

C++

Запускаємо пошук в глибину з вершини $1$ . Вершини графа нумеруються від $1$ до $n$ .

dfs(1);

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

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

Java реалізація

import java.util.*;
import java.io.*;

class Edge 
{
  int to;
  int weight;
  public Edge(int to, int weight) 
  {
    this.to = to;
    this.weight = weight;
  }
}

public class Main
{
  static int MOD = 1000000000;
  static ArrayList<Edge>[] g;	
  static int n, m;
  static long res;
  
  static int dfs(int v, int p)
  {
    int cnt = 1;
    for(Edge edge : g[v])
    {
      int to = edge.to;
      int w = edge.weight;
      if (to != p)
      {
        int c = dfs(to, v);
        res = (res + 1L * c * (n - c) * w) % MOD;
        cnt += c;
      }
    }
    return cnt;
  }
  
  public static void main(String[] args)
  {
    FastScanner con = new FastScanner(System.in);
    n = con.nextInt();
    g = new ArrayList[n+1]; // unchecked

    for (int i = 0; i <= n; i++)
      g[i] = new ArrayList<Edge>();

    for (int i = 1; i < n; i++)
    {
      int u = con.nextInt();
      int v = con.nextInt();
      int d = con.nextInt();
      g[u].add(new Edge(v,d));
      g[v].add(new Edge(u,d));
    }
    
    dfs(1,-1);
    System.out.println(res);    
  }
}   

class FastScanner 
{
  BufferedReader br;
  StringTokenizer st;

  public FastScanner(InputStream inputStream) 
  {
    br = new BufferedReader(new InputStreamReader(inputStream));
    st = new StringTokenizer("");
  }

  public String next() 
  {
    while (!st.hasMoreTokens()) 
    {
      try 
      {
        st = new StringTokenizer(br.readLine());
      } catch (Exception e) 
      {
        return null;
      }
    }
    return st.nextToken();
  }
 
  public int nextInt() 
  {
    return Integer.parseInt(next());
  }
}

Java

Python реалізація

Оголосимо значення модуля $M O D$ .

MOD = 1000000000

def dfs(v, p = -1):
  global res
  cnt = 1
  for to, w in g[v]:

Таким чином, внесок ребра $(v, t o)$ (v, to) у сумарну довжину всіх шляхів дорівнює

c \cdot (n - c) \cdot w

c \cdot (n - c) \cdot w

    if to != p:
      c = dfs(to, v)
      res = (res + c * (n - c) * w) % MOD
      cnt += c
  return cnt

Основна частина програми. Читаємо вхідні дані.

n = int(input())
g = [[] for _ in range(n + 1)]
res = 0

for _ in range(n - 1):
  u, v, d = map(int, input().split())
  g[u].append((v, d))
  g[v].append((u, d))

Python

Запускаємо пошук в глибину з вершини $1$ . Вершини графа нумеруються від $1$ до $n$ .

dfs(1)

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

print(res)