Advertisement
newb_ie

Segment Tree Basics

Apr 22nd, 2021
273
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 4.15 KB | None | 0 0
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. const int maxN = 100010;
  5. int n;
  6. int64_t t[4 * maxN];
  7. pair<int,int> minimum[4 * maxN];
  8. pair<int,int> maximum[4 * maxN];
  9. int ans = -1;
  10.  
  11. pair<int,int> combine (pair<int,int> a,pair<int,int> b) {
  12.     if (a.first < b.first) return a;
  13.     if (b.first < a.first) return b;
  14.     return make_pair(a.first,b.second + a.second);
  15. }
  16.  
  17. pair<int,int> combine_maximum(pair<int,int> a,pair<int,int> b) {
  18.     if (a.first > b.first) return a;
  19.     if (b.first > a.first) return b;
  20.     return make_pair(a.first,a.second + b.second);
  21. }
  22.  
  23. void build (int64_t a[],int v,int tl,int tr) {
  24.     if (tl == tr) {
  25.         t[v] = a[tl];
  26.         minimum[v] = make_pair(a[tl],1);
  27.         maximum[v] = make_pair(a[tl],1);
  28.     } else {
  29.         int tm = tl + (tr - tl) / 2;
  30.         build(a,v * 2,tl,tm);
  31.         build(a,v * 2 + 1,tm + 1,tr);
  32.         t[v] = t[v * 2] + t[v * 2 + 1];
  33.         minimum[v] = combine(minimum[2 * v],minimum[2 * v + 1]);
  34.         maximum[v] = combine_maximum(maximum[2 * v],maximum[2 * v + 1]);
  35.     }
  36. }
  37.  
  38. void point_update (int v,int tl,int tr,int pos,int new_val) {
  39.     if (tl == tr) {
  40.         t[v] = new_val;
  41.         minimum[v] = make_pair(new_val,1);
  42.         maximum[v] = make_pair(new_val,1);
  43.     } else {
  44.         int tm = tl + (tr - tl) / 2;
  45.         if (pos <= tm) {
  46.             point_update(v * 2,tl,tm,pos,new_val);
  47.         } else {
  48.             point_update(v * 2 + 1,tm + 1,tr,pos,new_val);
  49.         }
  50.         t[v] = t[v * 2] + t[v * 2 + 1];
  51.         minimum[v] = combine(minimum[v * 2],minimum[v * 2 + 1]);
  52.         maximum[v] = combine_maximum(maximum[2 * v],maximum[2 * v + 1]);
  53.     }
  54. }
  55.  
  56. int64_t sum_query (int v,int tl,int tr,int l,int r) {
  57.     if (l > r) return 0;
  58.     if (l == tl and r == tr) return t[v];
  59.     int tm = tl + (tr - tl) / 2;
  60.     return sum_query(v * 2,tl,tm,l,min(r,tm)) + sum_query(v * 2 + 1,tm + 1,tr,max(l,tm + 1),r);
  61. }
  62.  
  63. pair<int,int> minimum_query (int v,int tl,int tr,int l,int r) {
  64.     if (l > r) return make_pair(INT_MAX,0);
  65.     if (l == tl and r == tr) return minimum[v];
  66.     int tm = tl + (tr - tl) / 2;
  67.     return combine(minimum_query(v * 2,tl,tm,l,min(r,tm)),minimum_query(v * 2 + 1,tm + 1,tr,max(l,tm + 1),r));
  68. }
  69.  
  70. pair<int,int> maximum_query (int v,int tl,int tr,int l,int r) {
  71.     if (l > r) return make_pair(INT_MIN,0);
  72.     if (l == tl and r == tr) return maximum[v];
  73.     int tm = tl + (tr - tl) / 2;
  74.     return combine_maximum(maximum_query(v * 2,tl,tm,l,min(r,tm)),maximum_query(v * 2 + 1,tm + 1,tr,max(l,tm + 1),r));
  75. }
  76.  
  77. int find_kth_one_idx_query (int v,int tl,int tr,int k) {
  78.     if (k > t[v]) return -1; //if input[] contains less then k ones
  79.     if (tl == tr) return tl;
  80.     int tm = tl + (tr - tl) / 2;
  81.     if (k < t[v * 2]) {
  82.         return find_kth_one_idx_query(2 * v,tl,tm,k);
  83.     } else {
  84.         return find_kth_one_idx_query(2 * v + 1,tm + 1,tr,k - t[2 * v]);
  85.     }
  86. }
  87.  
  88. int find_kth_idx (int v,int tl,int tr,int k) {
  89.     if (tl == tr) return tl;
  90.     int tm = tl + (tr - tl) / 2;
  91.     if (maximum[2 * v].first < k) {
  92.         return find_kth_idx(2 * v + 1,tm + 1,tr,k);
  93.     } else {
  94.         return find_kth_idx(2 * v,tl,tm,k);
  95.     }
  96. }
  97.  
  98. void find_kth_idx_greater_equal_to_given_idx (int v,int tl,int tr,int k,int l) {
  99.     if (maximum[v].first < k) return;
  100.     if (ans >= 0 or tr < l) return;
  101.     if (tl == tr) {
  102.         if (maximum[v].first >= k) ans = tl;
  103.         return;
  104.     }
  105.     int tm = tl + (tr - tl) / 2;
  106.     find_kth_idx_greater_equal_to_given_idx(v * 2,tl,tm,k,l);
  107.     find_kth_idx_greater_equal_to_given_idx(v * 2 + 1,tm + 1,tr,k,l);
  108. }
  109.  
  110. int main () {
  111.      ios::sync_with_stdio(false);
  112.      cin.tie(nullptr);
  113.      cout.tie(nullptr);
  114.      int T = 1;
  115.      //~ cin >> T;
  116.      for (int test_case = 1; test_case <= T; ++test_case) {
  117.          int q;
  118.          cin >> n >> q;
  119.          int64_t input[n];
  120.          for (int i = 0; i < n; ++i) cin >> input[i];
  121.          build(input,1,0,n - 1);
  122.          //for sum_query sum_query (1,0,n - 1,l,r)
  123.          
  124.          //for minimum / maximum query pair<int,int> res = minimum / maximum_query(1,0,n - 1,l,r);
  125.          
  126.          //res.first -> min on the given segement and res.second -> count of min/max on a given segment
  127.          
  128.          //to find kth one find_kth_one_idx_query(1,0,n - 1,k)
  129.          
  130.          //to find kth idx where input[idx] >= x -> find_kth_idx(1,0,n - 1,x)
  131.          
  132.          //to find kth idx where idx >= l and input[idx] >= x -> find_kth_idx_greater_equal_to_given_idx(1,0,n - 1,k,l)
  133.      }
  134.      //cerr << "Time elapsed :" << clock() * 1000.0 / CLOCKS_PER_SEC << " ms" << '\n';
  135. }
  136.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement