Update 2106.Maximum-Fruits-Harvested-After-at-Most-K-Steps_v1.cpp

Author: wisdompeak

Binary_Search/2106.Maximum-Fruits-Harvested-After-at-Most-K-Steps/2106.Maximum-Fruits-Harvested-After-at-Most-K-Steps_v1.cpp

@@ -4,61 +4,33 @@ class Solution {
     {
         int n = fruits.size();
         vector<int>pos(n);
-        vector<int>val(n);
-        for (int i=0; i<n; i++)
-        {
-            pos[i] = fruits[i][0];
-            val[i] = fruits[i][1];
-        }
         vector<int>presum(n);
-        int cur = 0;
         for (int i=0; i<n; i++)
         {
-            cur += val[i];
-            presum[i] = cur;
+            pos[i] = fruits[i][0];
+            presum[i] = (i==0?0:presum[i-1]) + fruits[i][1];
         }
-
+        
         int ret = 0;
-        int j = 0;
-        while (j<n && pos[j]<startPos) j++;        
-        for (int i=0; i<n && pos[i]<=startPos; i++)
+        int mid = upper_bound(pos.begin(), pos.end(), startPos) - pos.begin();
+        int j = n-1;
+        for (int i=mid-1; i>=0; i--)
         {            
-            while (j<n && (startPos-pos[i]+(pos[j]-startPos)*2) <= k)
-            {                
-                ret = max(ret, presum[j] - (i==0 ? 0:presum[i-1]));
-                j++;
-            }
+            while (pos[j]>=startPos && startPos-pos[i]+(pos[j]-startPos)*2 > k)
+                j--;            
+            if (pos[j]>=startPos) ret = max(ret, presum[j] - (i==0?0:presum[i-1]));
+            else if (startPos-pos[i] <= k) ret = max(ret, presum[j] - (i==0?0:presum[i-1]));
         }
-
-        j = n-1;
-        while (j>=0 && pos[j]>startPos) j--;
-        for (int i=n-1; i>=0 && pos[i]>=startPos; i--)
+        
+        mid = lower_bound(pos.begin(), pos.end(), startPos) - pos.begin();
+        j = 0;        
+        for (int i=mid; i<n; i++)
         {
-            while (j>=0 && (pos[i]-startPos+(startPos-pos[j])*2 <=k))
-            {
-                
-                ret = max(ret, presum[i] - (j==0?0:presum[j-1]));
-                j--;
-            }
-        }
-
-        int i = upper_bound(pos.begin(), pos.end(), startPos) - pos.begin()-1;
-        int sum = 0;
-        while (i>=0 && startPos-pos[i]<=k) 
-        {
-            sum += val[i];
-            i--;
-        }
-        ret = max(ret, sum);
-
-        i = lower_bound(pos.begin(), pos.end(), startPos) - pos.begin();
-        sum = 0;
-        while (i<n && pos[i]-startPos<=k) 
-        {
-            sum += val[i];
-            i++;
+            while (pos[j]<=startPos && pos[i]-startPos+(startPos-pos[j])*2 > k)
+                j++;
+            if (pos[j]<=startPos) ret = max(ret, presum[i] - (j==0?0:presum[j-1]));
+            else if (pos[i]-startPos <= k) ret = max(ret, presum[i] - (j==0?0:presum[j-1]));
         }
-        ret = max(ret, sum);
 
         return ret;
     }