Cod sursa(job #2594967)

Utilizator vladth11Vlad Haivas vladth11 Data 6 aprilie 2020 20:32:06
Problema Diametrul unui arbore Scor 0
Compilator cpp-64 Status done
Runda Arhiva educationala Marime 2.42 kb
#include <bits/stdc++.h>
#define debug(x) cerr << #x << " " << x << "\n"

using namespace std;
typedef long long ll ;
typedef pair < pair <int,int>, pair <int,int> > pii;

const ll N = 400001;

vector <int> v[N];
pii dp[N];

void update(ll x, ll node) {
    if(dp[node].first.first < dp[x].first.first + 1 ) {
        swap(dp[node].first, dp[node].second);
        dp[node].first.first = dp[x].first.first + 1;
        dp[node].first.second = 1;
    } else if(dp[node].first.first == dp[x].first.first + 1) {
        dp[node].first.second ++;
    } else if(dp[node].second.first < dp[x].first.first + 1) {
        ///To prove that number of node doesn't affect my dp
        dp[node].second.first = dp[x].first.first + 1;
        dp[node].second.second = 1;
    } else if(dp[node].second.first == dp[x].first.first + 1) {
        dp[node].second.second ++;
    }
}

void DFS(ll node, ll parent) {
    dp[node].first.first = dp[node].first.second = dp[node].second.first = dp[node].second.second = 1;
    for(auto x : v[node]) {
        if(x != parent) {
            DFS(x, node);
            update(x, node);
        }
    }
}

ll maxim = 0,cnt = 0;

void Rerooting(ll node, ll parent) {

    if(dp[node].first.first > maxim) {
        maxim = dp[node].first.first;
        cnt = dp[node].first.second;
    } else if(dp[node].first.first == maxim) {
        cnt+=dp[node].first.second;
    }
    for(auto x : v[node]) {
        if(x != parent) {
           // cerr << node << " -> " << x << "\n";
            //debug(dp[node].first.first);
           // debug(dp[x].first.first);
            pii a = dp[node], b = dp[x];
            if(dp[node].first.first == dp[x].first.first + 1) {
                dp[node].first.second--;
                if(dp[node].first.second == 0)
                    swap(dp[node].first, dp[node].second);
            }else if(dp[node].second.first == dp[x].first.first + 1) {
                dp[node].second.second--;

            }
            update(node, x);
            Rerooting(x, node);
            dp[node] = a;
            dp[x] = b;
        }
    }
}

int main() {
    ll n,i;
    //ifstream cin("darb.in");
    //ofstream cout("darb.out");
    cin >> n;
    for(i = 1; i < n; i++) {
        ll x,y;
        cin >> x >> y;
        v[x].push_back(y);
        v[y].push_back(x);
    }
    DFS(1,0);
    Rerooting(1,0);
    cout << maxim << "\n";
    return 0;
}