Cod sursa(job #2050252)

Utilizator prisonbreakMichael Scofield prisonbreak Data 28 octombrie 2017 02:39:40
Problema Diametrul unui arbore Scor 100
Compilator cpp Status done
Runda Arhiva educationala Marime 2.39 kb
#include <cstdio>
#include <algorithm>
#include <fstream>
#include <iostream>
#include <vector>
#include <queue>
#include <map>
#include <cstring>
#include <string>
#include <set>
#include <stack>
#define pb push_back

#define mp make_pair
#define f first
#define s second
#define ll long long

using namespace std;


class Tree {
    int N;
    vector<vector<int>> edges;
    vector<int> first_max, second_max, combined_distance;
    vector<int> degree;
  public:
    Tree(int _N) {
        N = _N;
        edges.assign(N, vector<int>());
        first_max.assign(N, 0);
        second_max.assign(N, 0);
        combined_distance.assign(N, 0);
        degree.assign(N, 0);
    }

    void add_edge(int x, int y) {
        x--; y--;
        edges[x].push_back(y);
        edges[y].push_back(x);
        degree[x] += 1;
        degree[y] += 1;
    }

    int solve() {
        queue <int> Q;
        for (int i = 0; i < N; ++i) {
            if (degree[i] == 1) {
                Q.push(i);
            }
        }
        int diameter = 0;
        while(!Q.empty()) {
            int vertex = Q.front(); Q.pop();
            degree[vertex] -= 1;

            for (auto neighbour: edges[vertex]) {
                if (degree[neighbour] > 0) {
                    degree[neighbour] -= 1;
                    int candidate = first_max[vertex] + 1;
                    if (candidate > first_max[neighbour]) {
                        second_max[neighbour] = first_max[neighbour];
                        first_max[neighbour] = candidate;
                    } else {
                        if (candidate > second_max[neighbour]) {
                            second_max[neighbour] = candidate;
                        }
                    }
                    combined_distance[neighbour] = max(combined_distance[neighbour],\
                            first_max[neighbour] + second_max[neighbour]);

                }
                if (degree[neighbour] == 1) {
                    Q.push(neighbour);
                }
            }
            diameter = max(diameter, combined_distance[vertex]);
            diameter = max(diameter, first_max[vertex]);
        }
        return diameter + 1;
    }
};
int main() {
  ifstream cin("darb.in");
  ofstream cout("darb.out");

  int N; cin >> N;
  Tree T(N);

  for (int i = 0; i < N-1; ++i) {
      int x, y; cin >> x >> y;
      T.add_edge(x, y);
  }
  cout << T.solve() << "\n";
  return 0;
}