fork download
copy 
  1. #include <bits/stdc++.h>
  2. #include<ext/pb_ds/assoc_container.hpp>
  3. #include<ext/pb_ds/tree_policy.hpp>
  4.  
  5. using namespace std;
  6. using namespace __gnu_pbds;
  7. // find_by_order, order_of_key
  8. typedef tree<int, null_type, less<int>, rb_tree_tag, tree_order_statistics_node_update> oset;
  9. typedef tree<int, null_type, less_equal<int>, rb_tree_tag, tree_order_statistics_node_update> omset;
  10. #define int long long
  11. #define endl '\n'
  12. #define pi pair<int,int>
  13. #define adjs(name, type, size) vector<vector<type>>name(size)
  14. #define adjpass(name, type) vector<vector<type>>&name
  15. #define rest(name, val) memset(name,val,sizeof(name))
  16. #define all(x) x.begin(),x.end()
  17. #define killua ios_base::sync_with_stdio(false), cin.tie(NULL), cout.tie(0)
  18. //changes in dir:
  19. int dx[] = {0, 0, 1, -1, -1, 1, -1, 1};
  20. int dy[] = {1, -1, 0, 0, 1, 1, -1, -1};
  21. int cases = 0;
  22.  
  23. /***إِلا أَنْ يَشَاءَ اللَّهُ وَاذْكُرْ رَبَّكَ إِذَا نَسِيتَ وَقُلْ عَسَى أَنْ يَهْدِيَنِي رَبِّي لأَقْرَبَ مِنْ هَذَا رَشَدًا ***/
  24. const int N = 1e5 + 5;
  25. int a[N], ans[N]; // 1 , -1
  26. vector<int> adj[N];
  27. int n, m;
  28.  
  29. map<int, int> dfs(int node, int p, int sum = 0) {
  30.     int num = 1;
  31.     map<int, int> temp;
  32.     for (auto child: adj[node]) {
  33.         if (child == p)
  34.             continue;
  35.         map<int, int> k = dfs(child, node, sum + a[node]);
  36.         if (k.size() > temp.size())
  37.             swap(k, temp);
  38.         for (auto [x, y]: k) {
  39.             int nx = x;
  40.             temp[nx] += y;
  41.         }
  42.     }
  43.     ans[node] += temp[sum];
  44.     temp[a[node] + sum]++;
  45.     return temp;
  46. }
  47.  
  48. void gon() {
  49.     cin >> n >> m;
  50.     for (int i = 0; i < n; i++)cin >> a[i], (a[i] == 0 ? a[i] = -1 : a[i] = 1);
  51.     for (int i = 0; i < n - 1; i++) {
  52.         int u, v;
  53.         cin >> u >> v;
  54.         u--;
  55.         v--;
  56.         adj[u].push_back(v);
  57.         adj[v].push_back(u);
  58.     }
  59.     dfs(0, -1);
  60.     while (m--) {
  61.         int k;
  62.         cin >> k;
  63.         k--;
  64.         cout << ans[k] << endl;
  65.     }
  66. }
  67.  
  68. int32_t main() {
  69. #ifndef ONLINE_JUDGE
  70.     freopen("Input.txt", "r", stdin);
  71.     freopen("Output.txt", "w", stdout);
  72. #endif
  73.     killua;
  74.     int t = 1;
  75.     if (cases) cin >> t;
  76.     while (t--) {
  77.         gon();
  78.     }
  79. }
Success #stdin #stdout 0.01s 6852KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/8CAazS
language:
C++14 (gcc 8.3)
created:
11 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!