CHESUB - Editorial

Setter & Editorialist: Aman Dwivedi
Tester: Riley Borgard

Easy-Medium

DP

PROBLEM

You are given two integers N and K, and an array A of N integers. You have to choose K disjoint non-empty subarrays such that the score is maximized.

The score is calculated as follows:

\mathrm{Score}= \sum_{i=1}^{K} \mathrm{Sum}[i] \cdot i,

where \mathrm{Sum}[i] denotes sum of elements of i-th subarray. By the i-th subarray, we mean the i-th one in the order from left to right.

Find the maximum score that can be achieved.

EXPLANATION:

Our goal is to maximize the score, let us see how that can be done for the various subtasks.

K=1

Since we only need to choose 1 non-empty subarray in this subtask. And our goal is to maximize the score, then the optimal choice to choose subarray is the one that has maximum sum among all the choices available.

Hence, this subtask of the problem is reduced to finding the maximum subarray sum such that the subarray is non-empty. This can be easily done by Kadane’s Algorithm.

Time Complexity

O(N) per test case

Solution
#include <bits/stdc++.h>
using namespace std;

#define int long long

void solve()
{
int n,k;
cin>>n>>k;

int arr[n];
for(int i=0;i<n;i++)
cin>>arr[i];

int ans=arr[0];
int curr=0;

for(int i=0;i<n;i++)
{
curr+=arr[i];
ans=max(ans,curr);
if(curr<0)
curr=0;
}

if(k==1)
cout<<ans<<"\n";
else
exit(0);
}

int32_t main(){

ios_base::sync_with_stdio(0);
cin.tie(0);
int t;
cin>>t;

while(t--)
{
solve();
}

return 0;
}



K=2

We only need to choose 2 non-empty subarrays in this subtask. If there are two subarrays say S_1 and S_2, then we can say that the index of the first element of the subarray S_2 is always greater than the index of the last element of the subarray S_1.

Suppose we are at index i of the array A, then we can easily find the first subarray from the prefix of the array [A_1, A_2,\dots, A_i] and the other subarray from the suffix of the array [A_{i+1}, A_{i+2},\dots, A_N]. Since our goal is to maximize the score, we will always try to maximize the subarray sum of both prefix and suffix.

Hence for every index i of an array, we will find the maximum subarray sum in prefix [A_1, A_2,\dots, A_i] and maximum subarray sum in suffix [A_{i+1}, A_{i+2},\dots, A_N]. To optimally find the maximum subarray sum for every index i in prefix and suffix we can use prefix and suffix array where the i-th element of prefix array stores the maximum subarray sum possible in [A_1, A_2,\dots, A_i] and similarly for suffix array.

To find the maximum subarray sum we can use Kadane’s Algorithm.

Time Complexity

O(N) per test case

Solution
#include <bits/stdc++.h>
using namespace std;

#define int long long

void solve()
{
int n,k;
cin>>n>>k;

int arr[n];
for(int i=0;i<n;i++)
cin>>arr[i];

int pre[n],suf[n];

int ans=arr[0];
int curr=0;

for(int i=0;i<n;i++)
{
curr+=arr[i];
ans=max(ans,curr);
pre[i]=ans;
if(curr<0)
curr=0;
}

ans=arr[n-1];
curr=0;

for(int i=n-1;i>=0;i--)
{
curr+=arr[i];
ans=max(ans,curr);
suf[i]=ans;
if(curr<0)
curr=0;
}

ans=INT_MIN;

for(int i=0;i<n-1;i++)
{
int temp=pre[i]+suf[i+1]*2;
ans=max(ans,temp);
}

if(k==2)
cout<<ans<<"\n";
else
exit(0);
}

int32_t main(){

ios_base::sync_with_stdio(0);
cin.tie(0);
int t;
cin>>t;

while(t--)
{
solve();
}

return 0;
}



Original Constraints

Since the value of K is large enough this time to break our greedy solution. Notice that when we are at some index i of an array we had several choices like:

• Should we include this element in the previous subarray?
• Should we never take this element in any of the K subarrays?
• Should we start our new subarray from this element?

All these types of queries in our mind give us big hint towards Dynamic Programming as well as the states of the DP.

Now, let’s move towards Dynamic Programming.

Let’s create our DP as (i,j) which is defined as follows:

DP[i][j] is defined as the maximum score that is possible by taking the first i index of the array A and having exactly j subarrays such that the jth subarray ends at index i .

Suppose we are at index i of an array and we had selected j subarrays till now. Now for this index, i we had choices as:

• Take this element in the j-th subarray, then our transition will be as:
DP[i][j]=DP[i-1[[j]+j*A[i]
• Start a new subarray from this index of an array. Then the transition will be as follows:
DP[i][j]=max(DP[i-1][j-1],DP[i-2][j-1],\dots,DP[1][j-1])+j*A[i]

As we can see that the when we want to start a new subarray from index i if an array then there are several choices and we need to find the maximum among all since our goal is to maximize the score.

To find the maximum, we need to (i-1) iterations for every state leading the time complexity to N^2*K which surely results in TLE.

However, we can do better by using the concept of prefix array. We can define prefix array as:

prefix[i][j]=max(DP[i][j],DP[i-1][j],\dots,DP[1][j])

Hence instead of doing (i-1) iterations, we can simply do it in O(1) by using the concept of prefix array leading us to time complexity of O(N*K) which helps to get AC.

All we need to do is now compute the value of the prefix array. It is quite clear that any column of prefix array will always be sorted in non-decreasing order. Suppose we computed the value of DP[i][j], then we can update our prefix array as:

prefix[i][j]=max(prefix[i-1][j],DP[i][j])

TIME COMPLEXITY:

O(N*K) per test case

SOLUTIONS:

Setter
#include <bits/stdc++.h>
using namespace std;

#define int long long

int inf = 1e18;

int dp[100010][110];
int prefix[100010][110];

void solve()
{
int n, k;
cin >> n >> k;
int a[n], b[k];

for(int i=0; i<n; i++)
cin >> a[i];
for(int j=0; j<k; j++)
b[j]=j+1;

for(int i=0; i<n; i++){
dp[i][0] = 0;
prefix[i][0] = 0;
for(int j=1; j<=k; j++){
dp[i][j] = -inf;
prefix[i][j] = -inf;
}
}

dp[0][1] = a[0]*b[0];
prefix[0][1] = dp[0][1];

for(int i=1; i<n; i++){
for(int j=1; j<=k; j++){
dp[i][j] = max(dp[i-1][j], prefix[i-1][j-1]) + a[i]*b[j-1];
dp[i][j] = max(dp[i][j], -inf);
prefix[i][j] = max(prefix[i-1][j], dp[i][j]);
}
}

int ans = -inf;
for(int i=0; i<n; i++){
ans = max(ans, dp[i][k]);
}

cout << ans << "\n";
}

int32_t main(){
// freopen("sub5.in","r",stdin);
// freopen("sub5.out","w",stdout);
ios_base::sync_with_stdio(0);
cin.tie(0);
int t;
cin>>t;

while(t--)
{
solve();
}

return 0;
}


Tester
#include <bits/stdc++.h>

#define ll long long
#define sz(x) ((int) (x).size())
#define all(x) (x).begin(), (x).end()
#define vi vector<int>
#define pii pair<int, int>
#define rep(i, a, b) for(int i = (a); i < (b); i++)
using namespace std;
template<typename T>
using minpq = priority_queue<T, vector<T>, greater<T>>;

void solve() {
int n, k;
cin >> n >> k;
vector<ll> a(n + 1);
rep(i, 1, n + 1) {
cin >> a[i];
a[i] += a[i - 1];
}
vector<ll> dp(k + 1, LLONG_MIN), mx(k + 1, LLONG_MIN);
dp[0] = 0;
mx[0] = 0;
rep(i, 1, n + 1) {
for(int j = k; j >= 1; j--) {
if(mx[j - 1] == LLONG_MIN) {
dp[j] = LLONG_MIN;
mx[j] = LLONG_MIN;
}else {
dp[j] = max(dp[j], mx[j - 1] + a[i] * j);
mx[j] = max(mx[j], dp[j] - a[i] * (j + 1));
}
}
mx[0] = max(mx[0], dp[0] - a[i]);
}
cout << dp[k] << '\n';
}

int main() {
ios::sync_with_stdio(false);
cin.tie(0);
int te;
cin >> te;
while(te--) solve();
}

33 Likes

Nice recurrence! If you have more such DP problems (where recurrence is not just looking at previous or previous 2 states), can you share them?

16 Likes

Try this MEXUM

7 Likes

Losers were trying to solve the problem on their own. Magically these people came up with same solution at the same time.
@codechef
https://www.codechef.com/viewsolution/47272785
Solution: 47272764 | CodeChef tried to get away from plagirism by unnecessary declaration of variables Liike seriously dude
Solution: 47272268 | CodeChef
Solution: 47272253 | CodeChef
Solution: 47272227 | CodeChef
Solution: 47271959 | CodeChef
Solution: 47271796 | CodeChef
Solution: 47271108 | CodeChef – this guy wat an awesome way to hide it mann
Solution: 47271188 | CodeChef
Solution: 47270908 | CodeChef
and many more …
If u look carefully at the submission time almost everyone submitted at the same time within 5 min :-

29 Likes

In editorial we have handled 2 cases ie

1. to include ith element in jth group
2. to start jth group from ith element
But what about when we don’t want to include ith element into any group how is that case handled ?
Can anyone explain ?
2 Likes

According to constraints solution with O(N* K *2) should pass the test easily but why my code is giving TLE.
Can u point out anything which is creating TLE.
https://www.codechef.com/viewsolution/47222889

2 Likes

That’s covered in the prefix recurrence. If suppose i-2’th state was max, then you have essentially skipped i-1’th state at i’th iteration.

2 Likes

The dp[i][j] definition takes the i^{th} element in j^{th} subarray. After forming dp, you would calculate \max\limits_{i=1}^{n} dp[i][k] and this would be the required answer.

5 Likes

I have notice quite often that u guys just set so tight time limit that if someone doesn’t code the way ur editorial is he’s gonna get TLE. Grow up and give some margin in the time limit, nevertheless who cares Codechef rating is a joke.

9 Likes

In last 2 minutes, my rank went from 650 to 725. In just 2 freaking minutes!
All of them were able to solve this dp in the last 2 minutes!

7 Likes
2 Likes

Exactly. I also wrote a recursive dp with O(n *k *2) bit it got tle for 2 cases only. It means logic was correct but time limit was too strict

3 Likes

codechef is becoming worst platform for cp as there r a lot of cheaters on this site.

12 Likes

How the fuck are these companies recruiting/shortlisting based on these contests :-

4 Likes

Are u kidding? The time limit should have been 2 seconds for this problem, not everyone uses for loop dp…
Recursion memoization solution even though correct and having same complexity as many codes that got AC’d, fails only because of such tight constraints. There should be some overhead not so high that a worse complexity solution passes but also not so low either, it is usually seen on codeforces where they give some room in bugaboos having dp tags.

12 Likes

Exactly the same thing happened to me.
https://www.codechef.com/viewsolution/47282097

3 Likes

I was also getting TLE for last 2 subtasks and wasted my 1.5 hour thinking what further optimization this question requires!

2 Likes

Can someone please explain how this recursive solution works and what are the dp states in this solution , I don’t usually write iterative so that’s why I want to understand this solution , but having problem especially with the third state

#include <bits/stdc++.h>
using namespace std;
#define int long long

#pragma GCC target (“avx2”)
#pragma GCC optimization (“O3”)
#pragma GCC optimization (“unroll-loops”)

/*

*/

const int INF = 1e18;
const int N = 1e5 + 2, K = 101, C = 2;
int cache[N][K][C];

int n, k, res;
vector a(N);

int go(int at, int grp, int take) {
// cerr << at << ’ ’ << grp << ’ ’ << take << ‘\n’;
if (grp == k + 1) {
return 0;
}
if (at == n) {
return -INF;
}

int mem = cache[at][grp][take];
if (mem != -INF) {
return mem;
}

if (take == 1) {
res = a[at] * grp + max({go(at + 1, grp, 1), go(at + 1, grp + 1, 0), go(at + 1, grp + 1, 1)});
} else {
res = max(go(at + 1, grp, 0), go(at + 1, grp, 1));
}
return cache[at][grp][take] = res;


}

void run_case() {
cin >> n >> k;
for (int i = 0; i < n; i++) {
cin >> a[i];
}

for (int i = 0; i <= n; i++) {
for (int j = 0; j <= k; j++) {
cache[i][j][0] = -INF;
cache[i][j][1] = -INF;
}
}

cout << max(go(0, 1, 0), go(0, 1, 1)) << '\n';


}

signed main() {
ios_base::sync_with_stdio(false); cin.tie(nullptr);

int T = 1;
cin >> T;
for (int t = 1; t <= T; t++) {
// cout << "Case #" << t << ": ";
run_case();
}

return 0;


}

Check this
my recursive with memorization solution with comments
solution
I hope it would be helpful.

1 Like

so even the subarrays do not have the weight is it the same dp or is there any short answer possible. not having weight I mean instead of multiplying sum with 1,2,3 we will just have sum of sums