POI Task Meteors (49pts - TLE)

Jul 19th, 2022
785
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
1. #include <iostream>
2. #include <vector>
3. #include <tuple>
4. #include <algorithm>
5.
6. using namespace std;
7.
8. struct Node {
9.     long long sum = 0, prop = 0;
10.     Node * left, * right;
11. };
12. void propagate(Node * &base) {
13.     if (!base->left) base->left = new Node();
14.     if (!base->right) base->right = new Node();
15.     base->left->sum += base->prop;
16.     base->right->sum += base->prop;
17.     base->left->prop += base->prop;
18.     base->right->prop += base->prop;
19.     base->prop = 0;
20. }
21. void update(Node * &base, long long val, int qL, int qR, int l, int r) {
22.     if (qR <= l || r <= qL) return;
23.     if (qL <= l && r <= qR) {
24.         base->sum += (r-l)*val;
25.         base->prop += val;
26.         return;
27.     }
28.
29.     propagate(base);
30.     int m = (l+r)/2;
31.     update(base->left, val, qL, qR, l, m);
32.     update(base->right, val, qL, qR, m, r);
33.
34.     base->sum = base->left->sum + base->right->sum;
35. }
36. long long query(Node * &base, int index, int l, int r) {
37.     if (r-l==1) return base->sum;
38.     propagate(base);
39.     int m = (l+r)/2;
40.     if (index < m) return query(base->left, index, l, m);
41.     else return query(base->right, index, m, r);
42. }
43. void deleteTree(Node * &base) {
44.     if (base) {
45.         deleteTree(base->left);
46.         deleteTree(base->right);
47.         delete base;
48.     }
49. }
50.
51. int main() {
52.     ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
53.     int n, m; cin >> n >> m;
54.     vector<vector<int>> stateIndices(n);
55.     for (int i = 0; i < m; ++i) {
56.         int state; cin >> state;
57.         stateIndices[state-1].push_back(i);
58.     }
59.
60.     vector<int> minMeteors(n);
61.     for (int &x : minMeteors) cin >> x;
62.
63.     int k; cin >> k;
64.     vector<tuple<int, int, int>> meteors(k);
65.     for (auto &met : meteors) {
66.         cin >> get<0>(met) >> get<1>(met) >> get<2>(met);
67.         --get<0>(met), --get<1>(met);
68.     }
69.
70.     // Parallel Binary Search on number of meteors for range
71.     vector<int> stateL(n, 0), stateR(n, k+1);
72.     vector<pair<int, int>> q;
73.     while (true) {
74.         bool going = false; // at least one binary search still active
75.         q.clear();
76.         for (int s = 0; s < n; ++s) {
77.             if (stateL[s]+1 < stateR[s]) {
78.                 q.emplace_back((stateL[s]+stateR[s])/2, s);
79.                 going = true;
80.             }
81.         }
82.         sort(q.begin(), q.end());
83.
84.         if (!going) break;
85.
86.         int nextQ = 0; // index in q
87.
88.         // Build segment tree and ask queries along the way
89.         Node * root = new Node();
90.         auto updateTree = [&root, &m](int l, int r, int a) {
91.             update(root, a, l, r+1, 0, m);
92.         };
93.         for (int i = 0; i < k; ++i) {
94.             // Update the segment tree
95.             int qL = get<0>(meteors[i]), qR = get<1>(meteors[i]), a = get<2>(meteors[i]);
96.             if (qL <= qR) updateTree(qL, qR, a);
97.             else updateTree(qL, m-1, a), updateTree(0, qR, a);
98.
99.             // Make any state queries necessary
100.             while (nextQ < (int)q.size() && q[nextQ].first == i+1) {
101.                 int s = q[nextQ].second;
102.                 long long total = 0;
103.                 for (int index : stateIndices[s]) {
104.                     total += query(root, index, 0, m);
105.                 }
106.                 if (total < minMeteors[s]) stateL[s] = i+1;
107.                 else stateR[s] = i+1;
108.                 nextQ++;
109.             }
110.         }
111.         deleteTree(root);
112.     }
113.
114.     for (int s = 0; s < n; ++s) {
115.         if (stateR[s] == k+1) cout << "NIE\n";
116.         else cout << stateR[s] << '\n';
117.     }
118.     return 0;
119. }