Update Lagrange.mdx

Author: A1exL1ang

content/6_Advanced/Lagrange.mdx

@@ -120,41 +120,42 @@ using namespace std;
 
 #define ll long long
 
-const ll INF = (ll)300000 * 1000000000;
-
-int main() {
-	int n, k;
-	cin >> n >> k;
-
-	int a[n + 1];
-	for (int i = 1; i <= n; i++) { cin >> a[i]; }
-
-	auto solve_lambda = [&](ll lmb) {
-		pair<ll, ll> dp[n + 1][2];
-
-		dp[0][0] = {0, 0};
-		dp[0][1] = {-INF, 0};
-
-		for (int i = 1; i <= n; i++) {
-			dp[i][0] = max(dp[i - 1][0], dp[i - 1][1]);
-
-			dp[i][1] =
-			    max(make_pair(dp[i - 1][0].first + a[i] - lmb,
-			                  dp[i - 1][0].second + 1),
-			        make_pair(dp[i - 1][1].first + a[i], dp[i - 1][1].second));
-		}
-
-		return max(dp[n][0], dp[n][1]);
-	};
-
-	ll lo = 0;
-	ll hi = INF;
-	while (lo < hi) {
-		ll mid = (lo + hi) / 2;
-		solve_lambda(mid).second <= k ? hi = mid : lo = mid + 1;
-	}
-
-	cout << solve_lambda(lo).first + lo * k << endl;
+int main(){
+    int n, k;
+    cin >> n >> k;
+
+    int a[n];
+    for (int &i : a) cin >> i;
+
+    /*
+     * Returns {maximum sum, subarrays used (max if ties)} 
+     * if creating a subarray penalizes the sum by "lmb" and 
+     * there is no limit on the number of subarrays you can create
+    */
+    auto solveLambda = [&](ll lmb){
+        pair<ll, ll> dp[n][2];
+
+        dp[0][0] = {0, 0}; 
+        dp[0][1] = {a[0] - lmb, 1};
+    
+        for (int i = 1; i < n; i++){
+            dp[i][0] = max(dp[i - 1][0], dp[i - 1][1]);
+
+            dp[i][1] = max(
+                make_pair(dp[i - 1][0].first + a[i] - lmb, dp[i - 1][0].second + 1), 
+                make_pair(dp[i - 1][1].first + a[i], dp[i - 1][1].second)
+            );
+        }
+        return max(dp[n - 1][0], dp[n - 1][1]);
+    };
+    
+    ll lo = 0, hi = 1e18;
+
+    while (lo < hi){
+        ll mid = (lo + hi) / 2;
+        solveLambda(mid).second <= k ? hi = mid : lo = mid + 1;
+    }    
+    cout << solveLambda(lo).first + lo * k << endl;
 }
 ```