Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include <bits/stdc++.h> // Includes all standard libraries
- using namespace std;
- // Using long long for integer type for larger values
- #define int long long
- /**
- * @brief Performs DFS to find the farthest node from a given start node.
- * @param u Current node
- * @param p Parent node
- * @param dist Distance array (dist[u] = distance from root)
- * @param adj Adjacency list
- * @param max_dist Reference to store the maximum distance found
- * @param farthest_node Reference to store the node at max_dist
- */
- void find_farthest_node_dfs(int u, int p, vector<int>& dist,
- const vector<vector<int>>& adj, int& max_dist,
- int& farthest_node) {
- dist[u] = dist[p] + 1;
- for (int v : adj[u]) {
- if (v != p) {
- find_farthest_node_dfs(v, u, dist, adj, max_dist, farthest_node);
- }
- }
- // Check if this node is farther than the current max
- if (dist[u] > max_dist) {
- max_dist = dist[u];
- farthest_node = u;
- }
- }
- /**
- * @brief Performs DFS to populate distances for subtrees branching off the diameter.
- * @param u Current node
- * @param p Parent node (on the diameter)
- * @param max_dist_to_any_node The output array being populated
- * @param adj Adjacency list
- */
- void populate_subtree_dists_dfs(int u, int p, vector<int>& max_dist_to_any_node,
- const vector<vector<int>>& adj) {
- // The distance for a subtree node is its distance from the diameter
- // plus the diameter node's max distance to an endpoint.
- max_dist_to_any_node[u] = max_dist_to_any_node[p] + 1;
- for (int v : adj[u]) {
- if (v != p) {
- populate_subtree_dists_dfs(v, u, max_dist_to_any_node, adj);
- }
- }
- }
- signed main() {
- // Optimize C++ streams for speed
- ios::sync_with_stdio(false);
- cin.tie(NULL);
- int n;
- cin >> n;
- if (n == 1) {
- cout << 0 << '\n';
- return 0;
- }
- vector<vector<int>> adj(n + 1);
- for (int i = 0; i < n - 1; i++) {
- int u, v;
- cin >> u >> v;
- adj[u].push_back(v);
- adj[v].push_back(u);
- }
- vector<int> dist_from_ep1(n + 1, 0);
- int max_dist = -1;
- int endpoint1 = 0;
- // Use node 0 as a virtual parent with distance -1
- dist_from_ep1[0] = -1;
- // 1. Find one endpoint (endpoint1) of a diameter
- // We start from node 1 arbitrarily
- find_farthest_node_dfs(1, 0, dist_from_ep1, adj, max_dist, endpoint1);
- // 2. Find the other endpoint (endpoint2) and the true diameter length
- // dist_from_ep1 will now store distances from endpoint1
- max_dist = -1;
- int endpoint2 = 0;
- find_farthest_node_dfs(endpoint1, 0, dist_from_ep1, adj, max_dist, endpoint2);
- int diameter_len = max_dist;
- // 3. Reconstruct the diameter path (from endpoint2 back to endpoint1)
- vector<int> diameter_path;
- queue<int> q;
- q.push(endpoint2);
- while (!q.empty()) {
- int u = q.front();
- diameter_path.push_back(u);
- q.pop();
- if (u == endpoint1) break; // Reached the other end
- for (int v : adj[u]) {
- // Follow the parent pointer (dist from ep1 decreases by 1)
- if (dist_from_ep1[v] == dist_from_ep1[u] - 1) {
- q.push(v);
- break; // Found the unique parent on the path
- }
- }
- }
- // 4. Calculate max distance for every node
- // This is the max distance from any node 'u' to the farthest node in the tree
- // (which will always be either endpoint1 or endpoint2).
- vector<int> max_dist_to_any_node(n + 1, 0);
- // Set distances for the endpoints
- max_dist_to_any_node[endpoint1] = diameter_len;
- max_dist_to_any_node[endpoint2] = diameter_len;
- int path_len = diameter_path.size();
- // Iterate over internal nodes of the diameter
- for (int i = 1; i < path_len - 1; i++) {
- int u_on_path = diameter_path[i];
- // Distance for a node on the path is its max distance to either endpoint
- int dist_to_ep1 = dist_from_ep1[u_on_path];
- int dist_to_ep2 = diameter_len - dist_to_ep1;
- max_dist_to_any_node[u_on_path] = max(dist_to_ep1, dist_to_ep2);
- // 5. Run DFS for all subtrees branching off this diameter node
- for (int v_subtree : adj[u_on_path]) {
- // Avoid traversing back onto the diameter
- if (v_subtree != diameter_path[i - 1] && v_subtree != diameter_path[i + 1]) {
- populate_subtree_dists_dfs(v_subtree, u_on_path, max_dist_to_any_node, adj);
- }
- }
- }
- // Print the result for all nodes
- for (int i = 1; i <= n; i++) {
- cout << max_dist_to_any_node[i] << ' ';
- }
- cout << '\n';
- return 0;
- }
Advertisement
Add Comment
Please, Sign In to add comment