Дерево! (Розбір)

Рішення

Назвемо вузол листком, якщо його ступінь $1$. Нехай $L$ буде загальною кількістю листків у дереві.

Закріпіть дерево на $1$. Для кожного вузла $v$ нехай $s[v]$ буде кількістю листків у піддереві $v$ (відносно цього кореня). Ми можемо обчислити $s[\cdot ]$ за один DFS:

Розглянемо будь-яке ребро $(p,v)$, де $p$ є батьком $v$. Якщо ми розрізаємо це ребро, з'являються два компоненти:

- Сторона $v$ має $s[v]$ листків, але ступінь $v$ зменшується на $1$. Вузол стає листком якщо і тільки якщо його ступінь була $2$, а листок перестає бути листком якщо і тільки якщо його ступінь була $1$. Тому кількість листків на стороні $v$ після розрізання дорівнює

- Сторона $p$ має L-s[v] листків, і $p$ втрачає одне інцидентне ребро, тому його ступінь також зменшується на $1$. Таким чином, кількість листків на стороні $p$ після розрізання дорівнює

Ребро є збалансованим тоді і тільки тоді, коли A=B. Ми можемо обчислити L=s[1] з першого DFS, а потім у другому DFS по ребрах батько-дитина перевірити рівність вище і підрахувати, скільки разів вона виконується.

Ескіз коректності

Розрізання $(p,v)$ змінює статус листка лише на кінцевих точках $p$ і $v$ (всі інші зберігають той же ступінь), тому коригування на +[\deg=2]-[\deg=1] є як необхідним, так і достатнім. Оскільки $s[v]$ і L-s[v] рахують листки строго всередині кожної сторони перед врахуванням змін ступеня кінцевих точок, A=B точно характеризує збалансовані ребра.

Складність $\mathcal{O}(n)$ на тест (два обходи DFS).

#include <bits/stdc++.h>
using namespace std;

#define int int64_t
using vi=vector<int>;

void solve(){
    int n;cin>>n;
    vector<vi> g(n+1);
    for(int i=0,u,v;i<n-1;i++){
        cin>>u>>v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    vector<int> s(n+1);
    vector<char> vis(n+1,0);

    auto tour= [&](auto self,int v,int p)->void{
        vis[v]=1;
        s[v]+= (int)g[v].size()==1;
        for(int u:g[v]) if(u!=p) self(self,u,v), s[v]+=s[u];
    };
    tour(tour,1,0);
    int L=s[1], ans=0;
    fill(vis.begin(),vis.end(),0);

    auto dfs= [&](auto self,int v,int p)->void{
        vis[v]=1;
        for(int u:g[v]) if(u!=p){
            int A = s[u] + ((int)g[u].size()==2) - ((int)g[u].size()==1);
            int B = L - s[u] + ((int)g[v].size()==2) - ((int)g[v].size()==1);
            ans += (A==B);
            self(self,u,v);
        }
    };
    dfs(dfs,1,0);
    cout<<ans<<"\n";
}

signed main(){
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int t; if(!(cin>>t)) return 0;
    while(t--) solve();
    return 0;
}