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

wisdompeak · web-flow · commit bb57dab4961e · 2021-12-19T23:52:16.000-08:00
diff --git a/Binary_Search/2106.Maximum-Fruits-Harvested-After-at-Most-K-Steps/2106.Maximum-Fruits-Harvested-After-at-Most-K-Steps_v2.cpp b/Binary_Search/2106.Maximum-Fruits-Harvested-After-at-Most-K-Steps/2106.Maximum-Fruits-Harvested-After-at-Most-K-Steps_v2.cpp
@@ -12,25 +12,24 @@ class Solution {
         }
         
         int ret = 0;
-        for (int i=0; i<n; i++)
+        int mid = lower_bound(pos.begin(), pos.end(), startPos) - pos.begin();
+        for (int i=mid; i<n; i++)
         {
-            if (pos[i]<=startPos)
-            {
-                double d = k - (startPos-pos[i]);
-                if (d<0) continue;
-                int j = upper_bound(pos.begin(), pos.end(), startPos+d/2) - pos.begin() - 1;
-                if (j>=0 && j<n)
-                    ret = max(ret, presum[j] - (i==0?0:presum[i-1]));
-            }
-            else
-            {
-                double d = k - (pos[i]-startPos);
-                if (d<0) continue;   
-                int j = lower_bound(pos.begin(), pos.end(), startPos-d/2) - pos.begin();
-                if (j>=0 && j<n)
-                    ret = max(ret, presum[i]- (j==0?0:presum[j-1]));
-            }
+            if (pos[i] - startPos > k) break;
+            double d = (k-(pos[i]-startPos))/2;
+            int j = lower_bound(pos.begin(), pos.end(), startPos-d) - pos.begin();
+            if (j>=0 && j<n) ret = max(ret, presum[i] - (j==0?0:presum[j-1]));
+        }
+
+        mid = upper_bound(pos.begin(), pos.end(), startPos) - pos.begin()-1;
+        for (int i=mid; i>=0; i--)
+        {     
+            if (startPos - pos[i] > k) break;
+            double d = (k-(startPos-pos[i]))/2;
+            int j = upper_bound(pos.begin(), pos.end(), startPos+d) - pos.begin() - 1;
+            if (j>=0 && j<n) ret = max(ret, presum[j] - (i==0?0:presum[i-1]));
         }
+        
         return ret;
     }
 };
