Update Lagrange.mdx

Author: A1exL1ang

content/6_Advanced/Lagrange.mdx

@@ -120,38 +120,39 @@ using namespace std;
 
 #define ll long long
 
-const int MAX = 300000;
 const ll INF = (ll)300000 * 1000000000;
 
-int n, k, A[MAX + 1];
-pair<ll, ll> dp[MAX + 1][2];
-
-pair<ll, ll> solveLambda(ll lmb) {
-	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]);
-}
-
-int main() {
-	cin >> n >> k;
-
-	for (int i = 1; i <= n; i++) cin >> A[i];
-
-	ll lo = 0, hi = INF;
-
-	while (lo < hi) {
-		ll M = (lo + hi) / 2;
-		solveLambda(M).second <= k ? hi = M : lo = M + 1;
-	}
-	cout << solveLambda(lo).first + lo * k << endl;
+int main(){
+    int n, k;
+    cin >> n >> k;
+
+    int a[n + 1];
+    for (int i = 1; i <= n; i++) cin >> a[i];
+
+    auto solveLambda = [&](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, hi = INF;
+
+    while (lo < hi){
+        ll mid = (lo + hi) / 2;
+        solveLambda(mid).second <= k ? hi = mid : lo = mid + 1;
+    }    
+    cout << solveLambda(lo).first + lo * k << endl;
 }
 ```