chheem_tapak_tam_tam

CSES_Tree_Distances_I

Oct 27th, 2025 (edited)
64
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 4.80 KB | None | 0 0
  1. #include <bits/stdc++.h> // Includes all standard libraries
  2.  
  3. using namespace std;
  4.  
  5. // Using long long for integer type for larger values
  6. #define int long long
  7.  
  8. /**
  9.  * @brief Performs DFS to find the farthest node from a given start node.
  10.  * @param u Current node
  11.  * @param p Parent node
  12.  * @param dist Distance array (dist[u] = distance from root)
  13.  * @param adj Adjacency list
  14.  * @param max_dist Reference to store the maximum distance found
  15.  * @param farthest_node Reference to store the node at max_dist
  16.  */
  17. void find_farthest_node_dfs(int u, int p, vector<int>& dist,
  18.                             const vector<vector<int>>& adj, int& max_dist,
  19.                             int& farthest_node) {
  20.     dist[u] = dist[p] + 1;
  21.     for (int v : adj[u]) {
  22.         if (v != p) {
  23.             find_farthest_node_dfs(v, u, dist, adj, max_dist, farthest_node);
  24.         }
  25.     }
  26.     // Check if this node is farther than the current max
  27.     if (dist[u] > max_dist) {
  28.         max_dist = dist[u];
  29.         farthest_node = u;
  30.     }
  31. }
  32.  
  33. /**
  34.  * @brief Performs DFS to populate distances for subtrees branching off the diameter.
  35.  * @param u Current node
  36.  * @param p Parent node (on the diameter)
  37.  * @param max_dist_to_any_node The output array being populated
  38.  * @param adj Adjacency list
  39.  */
  40. void populate_subtree_dists_dfs(int u, int p, vector<int>& max_dist_to_any_node,
  41.                                 const vector<vector<int>>& adj) {
  42.     // The distance for a subtree node is its distance from the diameter
  43.     // plus the diameter node's max distance to an endpoint.
  44.     max_dist_to_any_node[u] = max_dist_to_any_node[p] + 1;
  45.     for (int v : adj[u]) {
  46.         if (v != p) {
  47.             populate_subtree_dists_dfs(v, u, max_dist_to_any_node, adj);
  48.         }
  49.     }
  50. }
  51.  
  52. signed main() {
  53.     // Optimize C++ streams for speed
  54.     ios::sync_with_stdio(false);
  55.     cin.tie(NULL);
  56.  
  57.     int n;
  58.     cin >> n;
  59.  
  60.     if (n == 1) {
  61.         cout << 0 << '\n';
  62.         return 0;
  63.     }
  64.  
  65.     vector<vector<int>> adj(n + 1);
  66.     for (int i = 0; i < n - 1; i++) {
  67.         int u, v;
  68.         cin >> u >> v;
  69.         adj[u].push_back(v);
  70.         adj[v].push_back(u);
  71.     }
  72.  
  73.     vector<int> dist_from_ep1(n + 1, 0);
  74.     int max_dist = -1;
  75.     int endpoint1 = 0;
  76.  
  77.     // Use node 0 as a virtual parent with distance -1
  78.     dist_from_ep1[0] = -1;
  79.    
  80.     // 1. Find one endpoint (endpoint1) of a diameter
  81.     // We start from node 1 arbitrarily
  82.     find_farthest_node_dfs(1, 0, dist_from_ep1, adj, max_dist, endpoint1);
  83.  
  84.     // 2. Find the other endpoint (endpoint2) and the true diameter length
  85.     // dist_from_ep1 will now store distances from endpoint1
  86.     max_dist = -1;
  87.     int endpoint2 = 0;
  88.     find_farthest_node_dfs(endpoint1, 0, dist_from_ep1, adj, max_dist, endpoint2);
  89.  
  90.     int diameter_len = max_dist;
  91.  
  92.     // 3. Reconstruct the diameter path (from endpoint2 back to endpoint1)
  93.     vector<int> diameter_path;
  94.     queue<int> q;
  95.     q.push(endpoint2);
  96.  
  97.     while (!q.empty()) {
  98.         int u = q.front();
  99.         diameter_path.push_back(u);
  100.         q.pop();
  101.  
  102.         if (u == endpoint1) break; // Reached the other end
  103.  
  104.         for (int v : adj[u]) {
  105.             // Follow the parent pointer (dist from ep1 decreases by 1)
  106.             if (dist_from_ep1[v] == dist_from_ep1[u] - 1) {
  107.                 q.push(v);
  108.                 break; // Found the unique parent on the path
  109.             }
  110.         }
  111.     }
  112.  
  113.     // 4. Calculate max distance for every node
  114.     // This is the max distance from any node 'u' to the farthest node in the tree
  115.     // (which will always be either endpoint1 or endpoint2).
  116.     vector<int> max_dist_to_any_node(n + 1, 0);
  117.  
  118.     // Set distances for the endpoints
  119.     max_dist_to_any_node[endpoint1] = diameter_len;
  120.     max_dist_to_any_node[endpoint2] = diameter_len;
  121.  
  122.     int path_len = diameter_path.size();
  123.    
  124.     // Iterate over internal nodes of the diameter
  125.     for (int i = 1; i < path_len - 1; i++) {
  126.         int u_on_path = diameter_path[i];
  127.        
  128.         // Distance for a node on the path is its max distance to either endpoint
  129.         int dist_to_ep1 = dist_from_ep1[u_on_path];
  130.         int dist_to_ep2 = diameter_len - dist_to_ep1;
  131.         max_dist_to_any_node[u_on_path] = max(dist_to_ep1, dist_to_ep2);
  132.  
  133.         // 5. Run DFS for all subtrees branching off this diameter node
  134.         for (int v_subtree : adj[u_on_path]) {
  135.             // Avoid traversing back onto the diameter
  136.             if (v_subtree != diameter_path[i - 1] && v_subtree != diameter_path[i + 1]) {
  137.                 populate_subtree_dists_dfs(v_subtree, u_on_path, max_dist_to_any_node, adj);
  138.             }
  139.         }
  140.     }
  141.  
  142.     // Print the result for all nodes
  143.     for (int i = 1; i <= n; i++) {
  144.         cout << max_dist_to_any_node[i] << ' ';
  145.     }
  146.     cout << '\n';
  147.  
  148.     return 0;
  149. }
  150.  
Advertisement
Add Comment
Please, Sign In to add comment