update range queries with seg tree

CSES - CSES Problem Set/Range_Queries/Range_Update_Query.cpp

@@ -7,11 +7,53 @@ using namespace std;
 using ll = long long;
 
 
+vector<ll> construct(vector<ll> &input) {
+    int n = input.size();
+    vector<ll> output(2 * n);
+    for (int i{}; i < n; i++) {
+        output[i + n] = input[i];
+    }
+    for (int i = n - 1; i >= 1; i--) {
+        output[i] = output[2*i] + output[2*i+1];
+    }
+    return output;
+}
 
-void update_range(int a, int b, ll k, vector<ll> &dif) {
-    dif[a] += k;
-    if (b != dif.size() - 1) {
-        dif[b+1] -= k;
+ll query(int k, vector<ll> &tree, int n) {
+    k += n;
+    ll s = 0;
+    int a = n;
+    while (a <= k) {
+        if (a % 2 == 1) {
+            s += tree[a];
+            a++;
+        }
+        if (k % 2 == 0) {
+            s += tree[k];
+            k--;
+        }
+        a /= 2; k /= 2;
+    }
+    return s;
+}
+
+
+void update_range(int a, int b, ll k, vector<ll> &tree, int n) {
+    if (b != n - 1) {
+        b += n;
+        tree[b + 1] -= k;
+        b++;
+        b /= 2;
+        for (; b >= 1; b /= 2) {
+            tree[b] = tree[2 * b] + tree[2 * b + 1];
+        }
+    }
+
+    a += n;
+    tree[a] += k;
+    a /= 2;
+    for (; a  >= 1; a /= 2) {
+        tree[a] = tree[2*a] + tree[2*a + 1];
     }
 }
 
@@ -20,18 +62,22 @@ int main() {
     int n, q; cin >> n >> q;
     vector<ll> input(n);
     vector<ll> difference;
-    for (auto &e: input){
-        cin >> e;
-        if (difference.size()) difference.push_back(e - input.back());
-        else difference.push_back(e);
+    for (int i{}; i < n; i++) {
+        cin >> input[i];
+        if (difference.size()) difference.push_back(input[i] - input[i - 1]);
+        else difference.push_back(input[0]);
     }
-
+    auto seg = construct(difference);
     for (int i{}; i < q; i++) {
         int t; cin >> t;
         if (t == 1) {
             int a, b; ll u; cin >> a >> b >> u;
             a--; b--;
-            update_range(a, b, u, difference);
+            update_range(a, b, u, seg, n);
+        } else {
+            int k; cin >> k;
+            k--;
+            cout << query(k, seg, n) << endl;
         }
     }