fork download
copy 
  1. #include <iostream>
  2. #include <vector>
  3. #include <numeric>
  4.  
  5. // These helpers and global variables should be placed outside the main and solve functions.
  6. const int MAXN_SIZE = 100005;
  7. std::vector<int> adj[MAXN_SIZE];
  8. long long node_values_storage[MAXN_SIZE];
  9. int tin[MAXN_SIZE], tout[MAXN_SIZE];
  10. long long flat_values[MAXN_SIZE];
  11. int timer;
  12. long long bit[MAXN_SIZE];
  13. int N_for_bit;
  14.  
  15. void bit_add(int idx, long long delta) {
  16.     for (; idx <= N_for_bit; idx += idx & -idx) {
  17.         bit[idx] += delta;
  18.     }
  19. }
  20.  
  21. long long bit_query(int idx) {
  22.     long long sum = 0;
  23.     for (; idx > 0; idx -= idx & -idx) {
  24.         sum += bit[idx];
  25.     }
  26.     return sum;
  27. }
  28.  
  29. long long range_sum(int l, int r) {
  30.     if (l > r) return 0;
  31.     return bit_query(r) - bit_query(l - 1);
  32. }
  33.  
  34. void dfs(int u, int p) {
  35.     tin[u] = ++timer;
  36.     flat_values[timer] = node_values_storage[u];
  37.     for (int v : adj[u]) {
  38.         if (v != p) {
  39.             dfs(v, u);
  40.         }
  41.     }
  42.     tout[u] = timer;
  43. }
  44.  
  45. bool is_ancestor(int u, int v) {
  46.     return tin[u] <= tin[v] && tout[u] >= tout[v];
  47. }
  48.  
  49. // Paste the implementation logic into the provided solve function.
  50. 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) {
  51.     N_for_bit = N;
  52.  
  53.     for (int i = 1; i <= N; ++i) {
  54.         adj[i].clear();
  55.         bit[i] = 0;
  56.         node_values_storage[i] = A[i - 1];
  57.     }
  58.  
  59.     for (const auto& edge : edges) {
  60.         adj[edge[0]].push_back(edge[1]);
  61.         adj[edge[1]].push_back(edge[0]);
  62.     }
  63.  
  64.     timer = 0;
  65.     // Assuming the root is always 1 as is common in these problems.
  66.     dfs(1, 0);
  67.  
  68.     for (int i = 1; i <= N; ++i) {
  69.         bit_add(i, flat_values[i]);
  70.     }
  71.  
  72.     std::vector<long long> results;
  73.     results.reserve(Q);
  74.  
  75.     for (const auto& q : query) {
  76.         int u = q[0];
  77.         int v = q[1];
  78.         int x = q[2];
  79.  
  80.         // Temporarily swap values if u is not an ancestor of v
  81.         bool swapped = false;
  82.         if (!is_ancestor(u, v)) {
  83.             swapped = true;
  84.             long long val_u = node_values_storage[u];
  85.             long long val_v = node_values_storage[v];
  86.  
  87.             bit_add(tin[u], val_v - val_u);
  88.             bit_add(tin[v], val_u - val_v);
  89.         }
  90.  
  91.         // Calculate subtree sum for x
  92.         results.push_back(range_sum(tin[x], tout[x]));
  93.  
  94.         // Revert the swap to restore original state for the next query
  95.         if (swapped) {
  96.             long long val_u = node_values_storage[u];
  97.             long long val_v = node_values_storage[v];
  98.             bit_add(tin[u], val_u - val_v);
  99.             bit_add(tin[v], val_v - val_u);
  100.         }
  101.     }
  102.  
  103.     return results;
  104. }
  105.  
  106. // This is the main function for handling I/O, as seen in your environment.
  107. int main() {
  108.     std::ios_base::sync_with_stdio(false);
  109.     std::cin.tie(NULL);
  110.     int T;
  111.     std::cin >> T;
  112.     while (T--) {
  113.         int N;
  114.         std::cin >> N;
  115.         std::vector<int> A(N);
  116.         for (int i = 0; i < N; ++i) {
  117.             std::cin >> A[i];
  118.         }
  119.  
  120.         std::vector<std::vector<int>> edges(N - 1, std::vector<int>(2));
  121.         for (int i = 0; i < N - 1; ++i) {
  122.             std::cin >> edges[i][0] >> edges[i][1];
  123.         }
  124.  
  125.         int Q;
  126.         std::cin >> Q;
  127.         std::vector<std::vector<int>> query(Q, std::vector<int>(3));
  128.         for (int i = 0; i < Q; ++i) {
  129.             std::cin >> query[i][0] >> query[i][1] >> query[i][2];
  130.         }
  131.  
  132.         std::vector<long long> result = solve(N, A, edges, Q, query);
  133.         for (size_t i = 0; i < result.size(); ++i) {
  134.             std::cout << result[i] << (i == result.size() - 1 ? "" : " ");
  135.         }
  136.         std::cout << "\n";
  137.     }
  138.     return 0;
  139. }
Success #stdin #stdout 0.01s 8812KB
comments (?)
stdin
copy 
2
5
4 1 4 2 3
1 2
2 3
2 4
1 5
1
3 5 2
2 5
stdout
copy 
6
0
https://ideone.com/yH8NZ8
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!