fork download
copy 
  1. #include <iostream>
  2. #include <vector>
  3. #include <numeric>
  4.  
  5. // Using global variables for simplicity as is common in competitive programming contests.
  6. // These will be re-initialized for each test case.
  7. std::vector<std::vector<int>> adj;
  8. std::vector<long long> node_values;
  9. std::vector<int> in_time;
  10. std::vector<int> out_time;
  11. std::vector<long long> subtree_sum;
  12. int timer;
  13.  
  14. /**
  15.  * @brief Performs DFS from a node to compute subtree sums and discovery/finish times.
  16.  * @param u The current node (1-based).
  17.  * @param p The parent of the current node (1-based).
  18.  */
  19. void dfs(int u, int p) {
  20.     in_time[u] = ++timer;
  21.     // node_values is 0-indexed, while node u is 1-indexed.
  22.     subtree_sum[u] = node_values[u - 1];
  23.  
  24.     for (int v : adj[u]) {
  25.         if (v != p) {
  26.             dfs(v, u);
  27.             subtree_sum[u] += subtree_sum[v];
  28.         }
  29.     }
  30.     out_time[u] = ++timer;
  31. }
  32.  
  33. /**
  34.  * @brief Checks if node u is an ancestor of node v.
  35.  * @param u The potential ancestor node (1-based).
  36.  * @param v The potential descendant node (1-based).
  37.  * @return True if u is an ancestor of v, false otherwise.
  38.  */
  39. bool is_ancestor(int u, int v) {
  40.     if (u == 0 || v == 0) return false; // A fictional parent 0 is not a real node.
  41.     // A node is considered an ancestor of itself.
  42.     return in_time[u] <= in_time[v] && out_time[u] >= out_time[v];
  43. }
  44.  
  45. /**
  46.  * @brief Solves the "Swap subtree" problem for one test case.
  47.  * @param N The number of nodes.
  48.  * @param A The vector of node values.
  49.  * @param edges The vector of edges in the tree.
  50.  * @param Q The number of queries.
  51.  * @param query The vector of queries.
  52.  * @return A vector of long long integers representing the answers to the queries.
  53.  */
  54. std::vector<long long> solve(int N, std::vector<int> A, std::vector<std::vector<int>> edges, int Q, std::vector<std::vector<int>> query) {
  55.     // 1. Initialize data structures for the current test case.
  56.     adj.assign(N + 1, std::vector<int>());
  57.     node_values.assign(A.begin(), A.end());
  58.     in_time.assign(N + 1, 0);
  59.     out_time.assign(N + 1, 0);
  60.     subtree_sum.assign(N + 1, 0LL);
  61.     timer = 0;
  62.  
  63.     // 2. Build the adjacency list representation of the tree.
  64.     for (const auto& edge : edges) {
  65.         adj[edge[0]].push_back(edge[1]);
  66.         adj[edge[1]].push_back(edge[0]);
  67.     }
  68.  
  69.     // 3. Pre-computation using DFS from the root (node 1).
  70.     dfs(1, 0); // The tree is rooted at 1, with a dummy parent 0.
  71.  
  72.     std::vector<long long> results;
  73.     results.reserve(Q);
  74.  
  75.     // 4. Process each query.
  76.     for (const auto& q : query) {
  77.         int U = q[0];
  78.         int V = q[1];
  79.         int X = q[2];
  80.  
  81.         // Check the condition for not performing the swap.
  82.         if (is_ancestor(U, V) || is_ancestor(V, U)) {
  83.             // No swap needed. The answer is the pre-calculated subtree sum.
  84.             results.push_back(subtree_sum[X]);
  85.         } else {
  86.             // Swap is performed. Calculate the new sum for X's subtree.
  87.             long long current_ans = subtree_sum[X];
  88.  
  89.             // If X is a *proper* ancestor of U, its subtree composition changes.
  90.             if (is_ancestor(X, U) && X != U) {
  91.                 current_ans += subtree_sum[V] - subtree_sum[U];
  92.             }
  93.  
  94.             // If X is a *proper* ancestor of V, its subtree composition changes.
  95.             if (is_ancestor(X, V) && X != V) {
  96.                 current_ans += subtree_sum[U] - subtree_sum[V];
  97.             }
  98.  
  99.             // Note: If X is not a proper ancestor (e.g., X=U, or X is in U's subtree, or X is unrelated),
  100.             // the conditions above are false, and the answer remains subtree_sum[X], which is correct.
  101.             results.push_back(current_ans);
  102.         }
  103.     }
  104.     return results;
  105. }
  106.  
  107.  
  108. // Main function to handle input/output as per the problem description.
  109. int main() {
  110.     std::ios::sync_with_stdio(false);
  111.     std::cin.tie(nullptr);
  112.  
  113.     int T;
  114.     std::cin >> T;
  115.     while (T--) {
  116.         int N;
  117.         std::cin >> N;
  118.         std::vector<int> A(N);
  119.         for (int i = 0; i < N; ++i) {
  120.             std::cin >> A[i];
  121.         }
  122.         std::vector<std::vector<int>> edges(N - 1, std::vector<int>(2));
  123.         for (int i = 0; i < N - 1; ++i) {
  124.             std::cin >> edges[i][0] >> edges[i][1];
  125.         }
  126.         int Q;
  127.         std::cin >> Q;
  128.         std::vector<std::vector<int>> query(Q, std::vector<int>(3));
  129.         for (int i = 0; i < Q; ++i) {
  130.             std::cin >> query[i][0] >> query[i][1] >> query[i][2];
  131.         }
  132.  
  133.         // Call the solve function with the parsed input.
  134.         std::vector<long long> result = solve(N, A, edges, Q, query);
  135.  
  136.         // Print the results for the current test case.
  137.         for (size_t i = 0; i < result.size(); ++i) {
  138.             std::cout << result[i] << (i == result.size() - 1 ? "" : " ");
  139.         }
  140.         std::cout << "\n";
  141.     }
  142.     return 0;
  143. }
  144.  
Success #stdin #stdout 0s 5324KB
comments (?)
stdin
copy 
2
5
4 1 4 2 3
1 2
2 3
2 4
1 5
1 
3 5 2 
5
stdout
copy 
6
0
https://ideone.com/NLIsPo
language:
C++ (gcc 8.3)
created:
4 days ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!