2D segment tree

CSES - CSES Problem Set/Range_Queries/Forest_Queries.cpp

@@ -0,0 +1,81 @@
+//  Forest Queries
+
+#include<bits/stdc++.h>
+
+using namespace std;
+
+vector<int> get_1D_seg_tree(string row) {
+    int n = row.size();
+    vector<int> output(2 * n);
+    for (int i{}; i < n; i++) {
+        output[i + n] = (row[i] == '*');
+    }
+    for (int i = n - 1; i >= 1; i--) {
+        output[i] = output[2 * i] + output[2 * i + 1];
+    }
+    return output;
+}
+
+vector<int> merge_2_seg_trees(vector<int> &first, vector<int> &second) {
+    vector<int> output(second.size());
+    for (int i{}; i < second.size(); i++) {
+        output[i] = first[i] + second[i];
+    }
+    return output;
+}
+
+int query_1D(vector<int> &seg1D, int n, int x1, int x2) {
+    x1 += n;
+    x2 += n;
+    int s = 0;
+    while (x1 <= x2) {
+        if (x1 % 2 == 1) {
+            s += seg1D[x1];
+            x1++;
+        }
+        if (x2 % 2 == 0) {
+            s += seg1D[x2];
+            x2--;
+        }
+        x1 /= 2;
+        x2 /= 2;
+    }
+    return s;
+}
+
+int query(vector<vector<int>> &seg, int n, int y1, int x1, int y2, int x2) {
+    y1 += n; y2 += n;
+    int s = 0;
+    while (y1  <= y2) {
+        if (y1 % 2 == 1) {
+            s += query_1D(seg[y1], n, x1, x2);
+            y1++;
+        }
+        if (y2 % 2 == 0) {
+            s += query_1D(seg[y2], n, x1, x2);
+            y2--;
+        }
+        y2 /= 2;
+        y1 /= 2;
+    }
+    return s;
+}
+
+int main() {
+    int n, q;
+    cin >> n >> q;
+    vector<vector<int>> TwoDSeg(2 * n);
+    for (int i{}; i < n; i++) {
+        string in; cin >> in;
+        TwoDSeg[i + n] = get_1D_seg_tree(in);
+    }
+    for (int i = n - 1; i >= 1; i--) {
+        TwoDSeg[i] = merge_2_seg_trees(TwoDSeg[2 * i], TwoDSeg[2 * i + 1]);
+    }
+    for (int i{}; i < q; i++) {
+        int y1, x1, y2, x2; cin >> y1 >> x1 >> y2 >> x2;
+        y1--; x1--; y2--; x2--;
+        cout << query(TwoDSeg, n, y1, x1, y2, x2) << endl;
+    }
+    return 0;
+}