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

За допомогою пошуку в глибину обчислимо XOR суму між вершиною $1$ і всіма іншими вершинами. XOR суму між вершинами $1$ і $v$ занесемо в $x[v]$. Нехай $ones$ — кількість одиниць, а $zeroes$ — кількість нулів у масиві x~(ones + zeroes = n). Тоді відповідь на задачу буде числом ones \cdot zeroes.

Приклад

Розглянемо граф, наведений у прикладі.

Біля кожної вершини $v$ записано XOR суму $x[v]$ між $1$ і $v$. Якщо x[v] = x[u] для деяких вершин $v$ і $u$, то XOR сума між ними дорівнює $0$, таким чином вносячи вклад $0$ в загальну суму. XOR сума для кожної пари вершин $(v,u)$, для якої x[v] \ne x[u], дає вклад $1$ у загальну суму.

Отже, відповідь дорівнює кількості пар вершин $(v,u)$, для яких x[v] \ne x[u]. Це число дорівнює ones \cdot zeroes = 2 \cdot 3 = 6. Парами вершин, що дають $1$ в загальну суму, будуть: (1, 2), (1, 4), (2, 3), (2, 5), (3, 4), (4, 5).

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

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

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

Функція dfs реалізує пошук в глибину, який обчислює XOR суму $x[v]$ між вершинами $1$ і $v$. Поточна XOR сума між $1$ і $v$ дорівнює cur\_xor. Предком вершини $v$ є $p$.

void dfs(int v, int cur_xor, int p = -1)
{
  x[v] = cur_xor;
  for (auto z : g[v])
  {
    int to = z.first;
    int w = z.second;
    if (to != p) dfs(to, cur_xor ^ w, v);
  }
}

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

scanf("%d", &n);
g.resize(n + 1);
x.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});
}

Запускаємо пошук в глибину з вершини $1$.

dfs(1, 0, -1);

Обчислюємо кількість нулів $zeroes$ і одиниць $ones$ у масиві $x$.

ones = 0;
for (i = 1; i <= n; i++)
  if (x[i] == 1) ones++;
zeroes = n - ones;

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

printf("%lld\n", 1LL * ones * zeroes);

Python реалізація

Функція dfs реалізує пошук в глибину, який обчислює XOR суму $x[v]$ між вершинами $1$ і $v$. Поточна XOR сума між $1$ і $v$ дорівнює cur\_xor. Предком вершини $v$ є $p$.

def dfs(v, cur_xor = 0, p = -1):
  x[v] = cur_xor
  for to, w in g[v]:
    if to != p:
      dfs(to, cur_xor ^ w, v)

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

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

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

x = [0] * (n + 1)

Запускаємо пошук в глибину з вершини $1$.

dfs(1)

Обчислюємо кількість нулів $zeroes$ і одиниць $ones$ у масиві $x$.

ones = sum(x)
zeroes = n - ones

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

print(ones * zeroes)