Advertisement
Guest User

Untitled

a guest
Jan 20th, 2019
87
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 3.68 KB | None | 0 0
  1. #include <iostream>
  2. #include <vector>
  3. #include <algorithm>
  4. #include <cassert>
  5. #include <queue>
  6. #include <set>
  7. #include <map>
  8. #include <unordered_map>
  9.  
  10. using namespace std;
  11.  
  12. typedef vector<int> V;
  13. typedef unordered_map<int, int> Map;
  14.  
  15. void printv(V& v) {
  16.     for (int elem : v) {
  17.         cout << elem << " ";
  18.     }
  19.     cout << endl;
  20. }
  21.  
  22. bool comp(int left, int right) {
  23.     return left < right;
  24. }
  25.  
  26. bool comprev(int left, int right) {
  27.     return left > right;
  28. }
  29.  
  30. class Heap {
  31.     bool minheap;
  32.     vector<int> h;
  33. public:
  34.     Heap(bool minheap) : minheap(minheap) {};
  35.  
  36.     void push(int v) {
  37.         h.push_back(v);
  38.         if (minheap) {
  39.             push_heap(h.begin(), h.end(), comprev);
  40.         } else {
  41.             push_heap(h.begin(), h.end(), comp);
  42.         }
  43.     }
  44.  
  45.     int pop() {
  46.         assert(h.size() > 0);
  47.         if (minheap) {
  48.             pop_heap(h.begin(), h.end(), comprev);
  49.         } else {
  50.             pop_heap(h.begin(), h.end(), comp);
  51.         }
  52.         int v = h.back();
  53.         h.pop_back();
  54.         return v;
  55.     }
  56.  
  57.     int peek() {
  58.         assert(h.size() > 0);
  59.         return h[0];
  60.     }
  61.  
  62.     int size() {
  63.         return h.size();
  64.     }
  65. };
  66.  
  67. void cleanup(Heap& heap, Map& count) {
  68.     while (true) {
  69.         auto search = count.find(heap.peek());
  70.         assert(search != count.end());
  71.         if (search->second == 0) {
  72.             heap.pop();
  73.         } else {
  74.             return;
  75.         }
  76.     }
  77. }
  78.  
  79. void checked_change(int key, int delta, Map& count) {
  80.     auto search = count.find(key);
  81.     assert(search != count.end());
  82.     search->second += delta;
  83. }
  84.  
  85. void first(int v, V& path, Map& count, Heap& minheap, Heap& maxheap, V& temp, vector<V>& routes, int m, int k, set<int>& out) {
  86.     path.push_back(v);
  87.     int curtemp = temp[v];
  88.     minheap.push(curtemp);
  89.     maxheap.push(curtemp);
  90.     checked_change(curtemp, 1, count);
  91.  
  92.     if (path.size() > m) {
  93.         int todel = temp[path[path.size() - m - 1]];
  94.         checked_change(todel, -1, count);
  95.     }
  96.  
  97.     cleanup(minheap, count);
  98.     cleanup(maxheap, count);
  99.  
  100.     if (path.size() >= m) {
  101.         int maxt = maxheap.peek();
  102.         int mint = minheap.peek();
  103.         //cerr << "maxt " << maxt << " mint " << mint << endl;
  104.         assert(maxt >= mint);
  105.         int diff = maxt - mint;
  106.         if (diff <= k) {
  107.             out.insert(path[path.size() - m]);
  108.         }
  109.     }
  110. }
  111.  
  112. void second(int v, V& path, Map& count, Heap& minheap, Heap& maxheap, V& temp, vector<V>& routes, int m, int k, set<int>& out) {
  113.     int curtemp = temp[v];
  114.  
  115.     if (path.size() > m) {
  116.         int toadd = temp[path[path.size() - m - 1]];
  117.         checked_change(toadd, 1, count);
  118.         minheap.push(toadd);
  119.         maxheap.push(toadd);
  120.     }
  121.  
  122.     checked_change(curtemp, -1, count);
  123.     path.pop_back();
  124. }
  125.  
  126. void testcase() {
  127.     int n, m, k;
  128.  
  129.     cin >> n >> m >> k;
  130.  
  131.     V temp;
  132.     Map count;
  133.     for (int i = 0; i < n; i++) {
  134.         int t;
  135.         cin >> t;
  136.         auto search = count.find(t);
  137.         if (search == count.end()) {
  138.             count.insert(pair<int, int>(t, 0));
  139.         }
  140.         temp.push_back(t);
  141.     }
  142.  
  143.     vector<V> routes(n);
  144.     for (int i = 0; i < n - 1; i++) {
  145.         int u, v;
  146.         cin >> u >> v;
  147.         routes[u].push_back(v);
  148.     }
  149.  
  150.     Heap minheap(true), maxheap(false);
  151.     V path;
  152.     set<int> out;
  153.  
  154.     set<int> processed;
  155.     vector<int> dfs;
  156.     dfs.push_back(0);
  157.  
  158.     while (dfs.size() > 0) {
  159.         int v = dfs.back();
  160.         dfs.pop_back();
  161.  
  162.         if (processed.count(v)) {
  163.             second(v, path, count, minheap, maxheap, temp, routes, m, k, out);
  164.             continue;
  165.         }
  166.  
  167.         processed.insert(v);
  168.         dfs.push_back(v);
  169.         for (int next : routes[v]) {
  170.             dfs.push_back(next);
  171.         }
  172.  
  173.         first(v, path, count, minheap, maxheap, temp, routes, m, k, out);
  174.     }
  175.  
  176.     V outsorted(out.begin(), out.end());
  177.     sort(outsorted.begin(), outsorted.end());
  178.     if (outsorted.size() == 0) {
  179.         cout << "Abort mission" << endl;
  180.     } else {
  181.         printv(outsorted);
  182.     }
  183. }
  184.  
  185. int main() {
  186.     ios_base::sync_with_stdio(false);
  187.     int t;
  188.     cin >> t;
  189.     while (t) {
  190.         //cerr << "testcase" << endl;
  191.         t--;
  192.         testcase();
  193.     }
  194.     return 0;
  195. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement