Update 2031.Count-Subarrays-With-More-Ones-Than-Zeros.cpp

Author: wisdompeak

Segment_Tree/2031.Count-Subarrays-With-More-Ones-Than-Zeros/2031.Count-Subarrays-With-More-Ones-Than-Zeros.cpp

@@ -1,24 +1,29 @@
-const int MAX_N = 200003;
-using LL = long long;
-
-class Solution {        
-    int OFFSET = 100001;
-    long long bitArr[MAX_N];
-    long long nums[MAX_N];  // Note: nums is 1-index
+class BIT{
+    public:
+    int N;    
+    vector<long long>bitArr; // Note: all arrays are 1-index
+    vector<long long>nums;
     long long M = 1e9+7;
 
-    // increase nums[i] by delta  (1-index)
+    void init(int N)
+    {
+        this->N = N;
+        bitArr.resize(N+1);
+        nums.resize(N+1);        
+    }
+
+    // increase nums[i] by delta
     void updateDelta(int i, long long delta) {
         int idx = i;
-        while (idx <= MAX_N)
+        while (idx <= N)
         {
             bitArr[idx]+=delta;
             bitArr[idx] %= M;
             idx+=idx&(-idx);
         }
     }
 
-    // sum of a range nums[1:j] inclusively, 1-index
+    // sum of a range nums[1:j] inclusively
     long long queryPreSum(int idx){
         long long result = 0;
         while (idx){
@@ -32,23 +37,29 @@ class Solution {
     // sum of a range nums[i:j] inclusively
     long long sumRange(int i, int j) {    
         return queryPreSum(j)-queryPreSum(i-1);
-    }
-    
+    } 
+};
+
+using LL = long long;
+LL OFFSET = 1e5+10;
+LL M = 1e9+7;
+class Solution {            
 public:
     int subarraysWithMoreZerosThanOnes(vector<int>& nums) 
     {
-        cout<<nums.size()<<endl;
+        BIT bit;
+        bit.init(2*OFFSET);
         int diff = 0;
         LL ret = 0;
-        updateDelta(0+OFFSET, 1);
+        bit.updateDelta(0+OFFSET, 1);
 
         for (int i=0; i<nums.size(); i++)
         {
             if (nums[i]==1) diff++;
             else diff--;
 
-            ret = (ret + queryPreSum(diff-1+OFFSET)) % M;                    
-            updateDelta(diff+OFFSET, 1);
+            ret = (ret + bit.queryPreSum(diff-1+OFFSET)) % M;                    
+            bit.updateDelta(diff+OFFSET, 1);
         }        
         return ret;        
     }